Home Cannot patch sqlalchemy database engine
Reply: 0

Cannot patch sqlalchemy database engine

Phylth
1#
Phylth Published in 2017-12-07 16:13:49Z

I am using SQLAlchemy (note: not Flask_SQLAlchemy) for a python 3 project, and I'm trying to write tests for the database by patching the engine with a test engine that points to a test database (as opposed to the production database). In the past, I successfully patched Session, and had working tests, but I recently switched to using the "insert" method, which is executed using engine.execute(), as opposed to a context managed session scope which was invoked using with session_scope() as session:

So heres the setup: I'm using a db_session module to establish a common session to be used by all DB functions:

import sys
import os
import logging

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from database.util import managed_session
import config

logger = logging.getLogger('default')

dirname = os.path.dirname
sys.path.append(dirname(dirname(__file__)))

engine = create_engine(config.POSTGRES_URI)

Session = sessionmaker(bind=engine)

def session_scope():
    return managed_session(Session)

and then in the crud_function file we have a setup as follows:

import logging
import re
from collections import defaultdict

from sqlalchemy import desc
from sqlalchemy.exc import IntegrityError

from database.db_session import session_scope, engine, Session
from database.models import *
from database.util import windowed_query
from sqlalchemy.dialects.postgresql import insert
import pandas as pd

def store_twitter_user(unprotected_row):
    '''
    Creates a TwitterUser object from the given attributes, adds it to the session, and then commits it to the database.
    :param attributes:
    :return:
    '''
    row = defaultdict(lambda: None, unprotected_row)

    pg_upsert(TwitterUser, row)

def pg_upsert(model, row):
    '''Performs an UPDATE OR INSERT ON CONFLICT (Upsert), which is a special SQL command for Postgres dbs.
    More info here: http://docs.sqlalchemy.org/en/latest/dialects/postgresql.html#insert-on-conflict-upsert
    '''
    try:
        insert_stmt = insert(model.__table__).values(row)
        do_update_stmt = insert_stmt.on_conflict_do_update(constraint=model.__table__.primary_key, set_=row)
        engine.execute(do_update_stmt)
        logger.debug('New {} stored successfully!'.format(type(object)))
        return True

    except IntegrityError as e:
        if re.search('violates foreign key constraint .* Key \(author_id\)=\(\d+\) is not present in table', str(e.args)):
            # Sends exception to celery task which will retry the task for a certain number of times
            raise

    except Exception as err:
        logger.error('pg_upsert: An error occurred while trying to store the new {}: {}'.format(model.__mapper__, err))
        return False

database.models just contains a bunch of classes used to create DB models for SQLAlchemy, like as follows:

class User(Base):
    __tablename__ = 'users'
    id = Column(BigInteger, primary_key=True)
    name = Column(String())
    screen_name = Column(String())
    location = Column(String(), index=True)
    friends_count = Column(Integer)
    created_at = Column(DateTime)
    time_zone = Column(String())
    lang = Column(String())

Now here's the test file:

engine = create_engine(config.POSTGRES_TEST_URI)

class TestDBCRUD(unittest.TestCase):
    ClassIsSetup = False
    ReadyForTeardown = False

    def setUp(self):
        """
        Creates all the tables in the test DB
        :return:
        """
        if not self.ClassIsSetup:
            print("SETTING UP!!!")
            Base.metadata.create_all(engine)
            self.__class__.ClassIsSetup = True

    def tearDown(self):
        """
        Deletes all test DB data and destroys the tables after a test is finished
        :return:
        """
        if self.ReadyForTeardown:
            print("TEARING DOWN!!!")
            Base.metadata.drop_all(engine)
            self.__class__.ReadyForTeardown = False

    @patch('database.crud.db_crud_functions.Session')
    @patch('database.crud.db_crud_functions.engine', autospec=True)
    @patch('database.db_session.Session')
    @patch('database.db_session.engine', autospec=True)
    def test_00_store_user(self, mock_session_engine, mock_session_Session, mock_engine, mock_session):
        print("testing store user!")
        Session = sessionmaker()
        Session.configure(bind=engine)
        mock_session_Session.return_value = Session()
        mock_session_engine.return_value = engine
        mock_session.return_value = Session()
        mock_engine.return_value = engine

        user = User(id=6789, screen_name="yeti")
        user_dict = {'id': 6789, 'screen_name': "yeti"}

        store_twitter_user(user_dict)

        with managed_session(Session) as session:
            retrieved_user = session.query(User).first()
            print("users are: {}".format(retrieved_user))
            self.assertEqual(user.id, retrieved_user.id)
            self.assertEqual(user.screen_name, retrieved_user.screen_name)

You'll notice a stupid amount of patches on top of the test function, and that is to show that I've tried to patch the engine and session from multiple locations. I've read that patches should be made where objects are used, and not where they are imported from, so I tried to cover all the bases. It doesn't matter, the test function always ends up inserting a user into the production database, and not into the test database. Then, when the retrieval happens, it returns None.

File "tests/testdatabase/test_db_crud_functions.py", line 59, in test_00_store_user
    self.assertEqual(user.id, retrieved_user.id)
AttributeError: 'NoneType' object has no attribute 'id'

Again, before pg_upsert was added, I used:

with session_scope() as session:
    session.add(something_here)

And session was successfully mocked to point to POSTGRES_TEST_URI, and not POSTGRES_URI. I'm at a loss here, please let me know if anything sticks out. Thanks!

You need to login account before you can post.

About| Privacy statement| Terms of Service| Advertising| Contact us| Help| Sitemap|
Processed in 0.314075 second(s) , Gzip On .

© 2016 Powered by mzan.com design MATCHINFO