Develop Reference

Unit testing a function that depends on database

I am running tests on some functions. I have a function that uses database queries. So, I have gone through the blogs and docs that say we have to make an in memory or test database to use such functions. Below is my function,
def already_exists(story_data,c):
# TODO(salmanhaseeb): Implement de-dupe functionality by checking if it already
# exists in the DB.
c.execute("""SELECT COUNT(*) from posts where post_id = ?""", (story_data.post_id,))
(number_of_rows,)=c.fetchone()
if number_of_rows > 0:
return True
return False
This function hits the production database. My question is that, when in testing, I create an in memory database and populate my values there, I will be querying that database (test DB). But I want to test my already_exists() function, after calling my already_exists function from test, my production db will be hit. How do I make my test DB hit while testing this function?

There are two routes you can take to address this problem:
Make an integration test instead of a unit test and just use a copy of the real database.
Provide a fake to the method instead of actual connection object.
Which one you should do depends on what you're trying to achieve.
If you want to test that the query itself works, then you should use an integration test. Full stop. The only way to make sure the query as intended is to run it with test data already in a copy of the database. Running it against a different database technology (e.g., running against SQLite when your production database in PostgreSQL) will not ensure that it works in production. Needing a copy of the database means you will need some automated deployment process for it that can be easily invoked against a separate database. You should have such an automated process, anyway, as it helps ensure that your deployments across environments are consistent, allows you to test them prior to release, and "documents" the process of upgrading the database. Standard solutions to this are migration tools written in your programming language like albemic or tools to execute raw SQL like yoyo or Flyway. You would need to invoke the deployment and fill it with test data prior to running the test, then run the test and assert the output you expect to be returned.
If you want to test the code around the query and not the query itself, then you can use a fake for the connection object. The most common solution to this is a mock. Mocks provide stand ins that can be configured to accept the function calls and inputs and return some output in place of the real object. This would allow you to test that the logic of the method works correctly, assuming that the query returns the results you expect. For your method, such a test might look something like this:
from unittest.mock import Mock
...
def test_already_exists_returns_true_for_positive_count():
mockConn = Mock(
execute=Mock(),
fetchone=Mock(return_value=(5,)),
)
story = Story(post_id=10) # Making some assumptions about what your object might look like.
result = already_exists(story, mockConn)
assert result
# Possibly assert calls on the mock. Value of these asserts is debatable.
mockConn.execute.assert_called("""SELECT COUNT(*) from posts where post_id = ?""", (story.post_id,))
mockConn.fetchone.assert_called()

The issue is ensuring that your code consistently uses the same database connection. Then you can set it once to whatever is appropriate for the current environment.
Rather than passing the database connection around from method to method, it might make more sense to make it a singleton.
def already_exists(story_data):
# Here `connection` is a singleton which returns the database connection.
connection.execute("""SELECT COUNT(*) from posts where post_id = ?""", (story_data.post_id,))
(number_of_rows,) = connection.fetchone()
if number_of_rows > 0:
return True
return False
Or make connection a method on each class and turn already_exists into a method. It should probably be a method regardless.
def already_exists(self):
# Here the connection is associated with the object.
self.connection.execute("""SELECT COUNT(*) from posts where post_id = ?""", (self.post_id,))
(number_of_rows,) = self.connection.fetchone()
if number_of_rows > 0:
return True
return False
But really you shouldn't be rolling this code yourself. Instead you should use an ORM such as SQLAlchemy which takes care of basic queries and connection management like this for you. It has a single connection, the "session".
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy_declarative import Address, Base, Person
engine = create_engine('sqlite:///sqlalchemy_example.db')
Base.metadata.bind = engine
DBSession = sessionmaker(bind=engine)
session = DBSession()
Then you use that to make queries. For example, it has an exists method.
session.query(Post.id).filter(q.exists()).scalar()
Using an ORM will greatly simplify your code. Here's a short tutorial for the basics, and a longer and more complete tutorial.

Use the SQLAlchemy session in custom Flask-SQLAlchemy query method

I would like to create a custom method to tell me whether a row exists in my database. The SQLAlchemy exists method returns a subquery but doesn't execute anything, so I wanted to create the method does_exist, which will simply return True or False. Here is my code:
from flask_sqlalchemy import SQLAlchemy, BaseQuery, Model
class CustomBaseQuery(BaseQuery):
def does_exist(self):
return db.session.query(self.exists()).scalar()
db = SQLAlchemy(app, query_class=CustomBaseQuery)
This actually does work, but it seems wrong to refer to db.session within the body of the method, which thus depends on later naming the SQLAlchemy instance db. I would like to find a way to reference the eventual db.session object in a more general way.
Full working example here: https://gist.github.com/wbruntra/3db7b630e6ffb86fe792e4ed5a7a9578

Though undocumented, the session used by the Query object is accessible as
self.session
so your more generic CustomBaseQuery could look like
class CustomBaseQuery(BaseQuery):
def does_exist(self):
return self.session.query(self.exists()).scalar()

Switching Collections and saving in Flask-Mongoengine

I am having trouble saving documents to a new collection and then deleting them from the old one.
I create a new object like so:
class Test(mongo.Document):
field = mongo.StringField()
t = Test(field="test")
t.switch_collection('default')
t.save()
t.switch_collection('switched')
t.save()
t.switch_collection('default')
t.delete()
It seems only to perform the first save to the default collection and then performs nothing after that. I have played around with a bunch of difference options such as reloading the object after every switch/save and from mongoengine context managers:
with switch_collection(Test, 'mongoswitch') as test:
test(field="switch").save()
My mongo settings look like (called first):
app.config["MONGODB_SETTINGS"] = {'db': 'TestDB'}
mongo = MongoEngine(app)
Using mongoengine 0.10 and pymongo 2.8.1 with Python 3.4 .
Anyone have an idea?
Much Thanks.

I see it's old question but maybe there's someone with the same problem.. I think it's because when you have a document that has id set, by calling t.save() you only update existing document in the collection. To really save it you need to call t.save(force_insert=True)

Are you getting any errors? It had worked for me. Check records in your mongo db collections.
One possible reason from mongoengine docs:
"Make sure any aliases have been registered with register_connection() or connect() before using the context manager."

Modify data as part of an alembic upgrade

I would like to modify some database data as part of an alembic upgrade.
I thought I could just add any code in the upgrade of my migration, but the following fails:
def upgrade():
### commands auto generated by Alembic - please adjust! ###
op.add_column('smsdelivery', sa.Column('sms_message_part_id', sa.Integer(), sa.ForeignKey('smsmessagepart.id'), nullable=True))
### end Alembic commands ###
from volunteer.models import DBSession, SmsDelivery, SmsMessagePart
for sms_delivery in DBSession.query(SmsDelivery).all():
message_part = DBSession.query(SmsMessagePart).filter(SmsMessagePart.message_id == sms_delivery.message_id).first()
if message_part is not None:
sms_delivery.sms_message_part = message_part
with the following error:
sqlalchemy.exc.UnboundExecutionError: Could not locate a bind configured on mapper Mapper|SmsDelivery|smsdelivery, SQL expression or this Session
I am not really understanding this error. How can I fix this or is doing operations like this not a possibility?

It is difficult to understand what exactly you are trying to achieve from the code excerpt your provided. But I'll try to guess. So the following answer will be based on my guess.
Line 4 - you import things (DBSession, SmsDelivery, SmsMessagePart) form your modules and then you are trying to operate with these objects like you do in your application.
The error shows that SmsDelivery is a mapper object - so it is pointing to some table. mapper objects should bind to valid sqlalchemy connection.
Which tells me that you skipped initialization of DB objects (connection and binding this connection to mapper objects) like you normally do in your application code.
DBSession looks like SQLAlchemy session object - it should have connection bind too.
Alembic already has connection ready and open - for making changes to db schema you are requesting with op.* methods.
So there should be way to get this connection.
According to Alembic manual op.get_bind() will return current Connection bind:
For full interaction with a connected database, use the “bind” available from the context:
from alembic import op
connection = op.get_bind()
So you may use this connection to run your queries into db.
PS. I would assume you wanted to perform some modifications to data in your table. You may try to formulate this modification into one update query. Alembic has special method for executing such changes - so you would not need to deal with connection.
alembic.operations.Operations.execute
execute(sql, execution_options=None)
Execute the given SQL using the current migration context.
In a SQL script context, the statement is emitted directly to the output stream. There is no return result, however, as this function is oriented towards generating a change script that can run in “offline” mode.
Parameters: sql – Any legal SQLAlchemy expression, including:
a string a sqlalchemy.sql.expression.text() construct.
a sqlalchemy.sql.expression.insert() construct.
a sqlalchemy.sql.expression.update(),
sqlalchemy.sql.expression.insert(), or
sqlalchemy.sql.expression.delete() construct. Pretty much anything
that’s “executable” as described in SQL Expression Language Tutorial.

Its worth noting that if you do this, you probably want to freeze a copy of your orm model inside the migration, like this:
class MyType(Base):
__tablename__ = 'existing_table'
__table_args__ = {'extend_existing': True}
id = Column(Integer, ...)
..
def upgrade():
Base.metadata.bind = op.get_bind()
for item in Session.query(MyType).all():
...
Otherwise you'll inevitably end up in a situation where you orm model changes and previous migrations no longer work.
Particularly note that you want to extend Base, not the base type itself (app.models.MyType) because your type might go away as some point, and once again, your migrations will fail.

You need to import Base also and then
Base.metatada.bind = op.get_bind()
and after this you can use your models like always without errors.

Django Unittesting adding a variable to the session

In my unittest I need to add a variable to the session, because that variable is used in the view which is being tested. The django documentation says this is possible in the following way (https://docs.djangoproject.com/en/1.3/topics/testing/#django.test.client.Client.session):
def test_something(self):
session = self.client.session
session['somekey'] = 'test'
session.save()
This code example actually doesn't work, because you will get the error that a dict doesn't have a function save. I also tried various other ways to change the contents inside the session dict, but haven't found a way to change it yet.
I know what it means, what i get back is a dict object so it doesnt have the save function. But the session dict also doesn't update when adding keys.
The documentation statues when using self.client.session it should return a SessionStore object instead of a dictionary.

Seems right now there is a bug in Django, so it doesn't work for unauthenticated users to change the session. This is the corresponding ticket: https://code.djangoproject.com/ticket/11475 .
A work around is to create a dummy view in which the session variables are set and calling that view with the Client.get('url_of_dummy_view').

The ticket referenced by #Sam Stoelinga has been closed as a duplicate. The corresponding ticket can be found here.
To save session data you will have to use an authenticated user. I found a solution here.
class BlogAdminTestCase(TestCase):
def setUp(self):
# Setup Test User
User.objects.create_user(
username='foo',
password='bar'
)
# Must login to modify session variables
self.client.login(username='foo', password='bar')
s = self.client.session
s['my_session_variable'] = 'Yay!'
s.save()