CODEWITHSUNDEEP
pythonsessionsqlalchemy
FranGoitia
FranGoitia

Reputation: 1993

SQLAlchemy mock Session connection

Everywhere I need to query the database in my application I import the Session from this file.

from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import scoped_session, sessionmaker

db_name = 'analytics'
db_url = 'postgresql+psycopg2://**:**@localhost:5432/{0}'.format(db_name)
Engine = create_engine(db_url, echo=False)
SessionMaker = sessionmaker(bind=Engine, autoflush=False)
Session = scoped_session(SessionMaker)

Base = declarative_base()

I'm writing some unit tests and need to mock the Session object so that whenever it is imported in my app, it connects to the unittest database rather than the analytics database.

class FacebookTest(unittest.TestCase):

    @classmethod
    def setUpClass(self):
        """
        create unittest database and insert Premier League 2015-2016 season
        """
        con = connect(user='**', host='localhost', password='**', database='postgres')
        con.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
        cur = con.cursor()
        cur.execute('DROP DATABASE IF EXISTS unittest')
        cur.execute('CREATE DATABASE unittest')
        cur.close()
        con.close()
        Engine = create_engine('postgresql+psycopg2://**:**@localhost:5432/unittest')
        Base.metadata.create_all(Engine)
        Session.connection = Engine.connect()

Upvotes: 0

Views: 4349

Answers (1)

Myke Bilyanskyy
Myke Bilyanskyy

Reputation: 678

Set your mock database on a different port, for example 5433 and change the port depending on which DB you want to connect to.

You can do like this. Let's say you have a db.py where you initiate your sessions and import them all over. Modify it like this (ignore reference to configurator, that's my internal module):

import psycopg2
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
import configurator

Base = declarative_base()

def connect():
    cfg = configurator.Configurator("database")
    return psycopg2.connect(
        user=cfg.user, host=cfg.host, dbname=cfg.db, port=cfg.port, password=cfg.pwd
    )

def connect_test():
    conn_string = (
        "host='localhost' user='postgres' password='yourPassword' port='5433'"
    )
    return psycopg2.connect(conn_string)

def Session(database='dev'):
    if database == 'dev':
        engine = create_engine("postgresql://", creator=connect)
    elif database == 'test':
        engine = create_engine("postgresql://", creator=connect_test)
    Session = sessionmaker(bind=engine)
    return Session()

UPDATE: Important correction - added brackets in return session(), sorry about that.

UPDATE2: Changed code after explicit testing.

Upvotes: 1

Related Questions