Can we use a different DB for running pytest? - python

So I am using FastAPI for creating my app server, I have some urls in it which I do some crud operation on database. For Eg: when I hit a url /createuser/ with a json body it inserts a record in db.
So I am using pytest to test this things but pytest seems to use the same db which is being used by my FastAPI app. I want to create a mock db so that it uses another db just for running tests. I have used pytest.fixtures as well but I guess something is wrong. Please help with this thing.
Also I am using postgres as my database.
Main App where FastAPI app is created
from mock_server.database import engine, SessionLocal, Base
def create_app(connection):
"""This function creates the FastAPI app server.
Returns:
app[FastAPI] -- The main app server
"""
Base.metadata.create_all(bind=connection)
print("from main:app",connection)
app = FastAPI(
title="Masquerader - Mock Server",
description="""This project allows you to mock any system,
or service that you wish to."""
)
return app
app = create_app(engine)
def get_db():
"""To Get the current db connection.
Yields:
[db] -- [The current connection]
"""
try:
db = SessionLocal()
yield db
finally:
db.close()
#app.post("/createuser/", response_model=schemas.User)
async def create_user(user: schemas.UserCreate, db: Session = Depends(get_db)):
"""Endpoint for creating a user.
Arguments:
user {schemas.UserCreate} -- JSON Body with user details to create
Keyword Arguments:
db {Session} -- Current db connection
"""
user = crud.create_user(db=db, user=user)
if user is None:
raise HTTPException(status_code=400, detail="User Already Exists")
return user
My Database class where db connections are defined
from config import current_config
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
# sqlalchemy_db_url = "postgresql://fyndlocal:fynd#123#localhost:5432/mockdb"
if os.environ.get("ENV") == "development":
engine = create_engine(current_config.POSTGRES_MASQUERADER)
if os.environ.get("ENV") is None:
print("Lol test hai yeh")
engine = create_engine(current_config.MASQUERADER_LOCAL)
if os.environ.get("ENV") == "pre-production":
os.environ.__setitem__("POSTGRES_USER", "runner")
os.environ.__setitem__("POSTGRES_PASS", "")
os.environ.__setitem__("POSTGRES_HOST", "postgres")
engine = create_engine(current_config.POSTGRES_TEST_GITLAB)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
My CRUD operation class
def create_user(db: Session, user: schemas.UserCreate):
"""Operation to create a user in the db.
Arguments:
db {Session} -- gets current db session
user {schemas.UserCreate} -- JSON Body that contains
user details to be created.
Returns:
db_user[dict] -- Details of created user.
"""
hashed_pass = common.hash_generate(user.password)
old_user = (
db.query(models.users).filter(models.users.name == user.name).first()
)
if old_user is not None:
return None
db_user = models.users(
name=user.name,
password=hashed_pass,
is_active=user.is_active,
is_admin=user.is_admin,
)
db.add(db_user)
db.commit()
db.refresh(db_user)
return db_user
My Test Class (in which I want it to use a mock db)
from main import app
import pytest
from sqlalchemy import create_engine
from starlette.testclient import TestClient
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
#pytest.yield_fixture(scope="module")
def application():
"""Yiest TestClient from FastAPI.
Yields:
app[TestClient] -- Testing based application
"""
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=cleanup_database)
db = SessionLocal()
yield TestClient(app)
#pytest.fixture(scope="module")
def cleanup_database():
"""Creates a mock database for testing purposes.
Creates a mock database on server for testing and deletes once done.
"""
username = os.environ.get("POSTGRES_USER", "fyndlocal")
password = os.environ.get("POSTGRES_PASS", "fynd#123")
postgres_host = os.environ.get("POSTGRES_HOST", "localhost")
postgres_port = os.environ.get("POSTGRES_PORT_5432_TCP_PORT", 5432)
postgres_db = os.environ.get("POSTGRES_DB", "mockdb")
if not password:
db_dict = {
"username": username,
"password": password,
"host": postgres_host,
"port": postgres_port,
"db": postgres_db,
}
default_db_url = current_config.POSTGRES_NOPASS_DSN.format(**db_dict)
print("if no pass", default_db_url)
else:
db_dict = {
"username": username,
"password": password,
"host": postgres_host,
"port": postgres_port,
"db": postgres_db,
}
default_db_url = current_config.POSTGRES_PASS_DSN.format(**db_dict)
print("if pass", default_db_url)
test_masqueraderdb_url = current_config.POSTGRES_TEST_DSN.format(**db_dict)
print("POSTGRES Config")
print(db_dict)
print(test_masqueraderdb_url)
db_name = test_masqueraderdb_url.split("/")[-1]
engine = create_engine(default_db_url)
conn = engine.connect()
# db.drop_all(engine)
# db.create_all(engine)
try:
conn.execution_options(isolation_level="AUTOCOMMIT").execute(
f"CREATE DATABASE {db_name}"
)
except Exception as e:
print("this stage", e)
rv = create_engine(test_masqueraderdb_url)
db.create_all(rv)
yield rv
db.drop_all(rv)
def test_create_user(application,cleanup_database):
"""Test to create user exists in db."""
response = application.post(
"/createuser/",
json={
"name": "test",
"is_active": True,
"is_admin": True,
"password": "test123",
},
)
expected_resp = {
"name": "test",
"is_active": True,
"is_admin": True,
"id": 1,
"urls": [],
}
assert response.json() == expected_resp
The following test is failing because it always gets the user present in my current db.
Thank you.

Give a try to overriding dependencies.
You should come up with a fixture similiar to the following (untested) code :
#pytest.fixture(scope="module")
def application(cleanup_database):
def get_mocked_db():
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=cleanup_database)
db = SessionLocal()
try:
yield db
finally:
db.close()
app = create_app(cleanup_database)
app.dependency_overrides[get_db] = get_mocked_db # see https://fastapi.tiangolo.com/advanced/testing-dependencies/
yield TestClient(app)
Note that the application fixture depends on the cleanup_database fixture.

Related

SQL Alchemy - Base.metadata.create_all(bind=engine) not creating tables in test db

I am using FastAPI to build an API backend for my URL shortener. I have the database connected to the API as a dependency. For the CRUD operations, I am using the SQL Alchemy ORM.
The code for my main app works perfectly fine and performs all the major CRUD operations I have mapped through the API endpoints.
The problem arises when I try to override the DB dependency to use a test db instead of my production db for testing purposes.
There are no errors associated with this override, however, the test database does not contain any of the tables that would be created when Base.metadata.creat_all(bind=engine) is called.
When running the tests using pytest, it gives me this error:
sqlalchemy.exc.ProgrammingError: (pymysql.err.ProgrammingError) (1146, "Table 'testurldb.urls' doesn't exist")
The code for my tests:
engine = create_engine(
"mysql+pymysql://{user}:{password}#{ip}:{port}/testurldb".format(
user=user, password=password, ip=ip, port=port
)
)
Session = sessionmaker(bind=engine)
Base.metadata.create_all(bind=engine)
def overrideDB():
db = Session()
try:
yield db
finally:
db.close()
app.dependency_overrides[get_db] = overrideDB
client = TestClient(app)
The module where Base is instantiated:
engine = create_engine(
"mysql+pymysql://{root}:{password}#{ip}:{port}/urldb".format(
root=root, password=password, ip=ip, port=port
)
)
SessionLocal = sessionmaker(bind=engine)
Base = declarative_base()
The table that extends Base:
class URL(Base):
__tablename__ = "urls"
short_url = Column(String(256), primary_key=True, unique=True, nullable=False)
long_url = Column(String(256), nullable=False, unique=True)
time = Column(String(256), nullable=False)
def __init__(self, short_url, long_url, time):
self.short_url = short_url
self.long_url = long_url
self.time = time
There seems to be nothing wrong with the imports, so I don't understand why it's not creating the tables.
And second, it might be useful information that my main production db already has the Tables created.
Create a file and import all models into it, at the end of all imports, put the Model import. When doing this, try to create the models again through Base.metadata
i hope it'll helps you
db.py
Base = declarative_base()
async def init_db():
try:
Base.metadata.create_all(bind=engine)
except Exception as e:
raise e
main.py
#app.on_event("startup")
async def on_startup():
await init_db()

FastAPI psycopg2 error when creating database for tests

I'm trying to test my database in my FastAPI app.
from fastapi.testclient import TestClient
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from app.main import app
from app.db import Base, get_db, get_session
SQLALCHEMY_DATABASE_URL = "sqlite:///./test.db"
engine = create_engine(
SQLALCHEMY_DATABASE_URL, connect_args={"check_same_thread": False}
)
TestingSessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base.metadata.create_all(bind=engine)
def override_get_db():
try:
db = TestingSessionLocal()
yield db
finally:
db.close()
app.dependency_overrides[get_db] = override_get_db
client = TestClient(app)
def test_get_users():
response = client.get("/users")
assert response.status_code == 200
This gives me an error sqlalchemy.exc.OperationalError: (psycopg2.OperationalError) could not translate host name "db" to address: Name or service not known.
In my app.db, this is how I define the engine:
engine = create_engine(settings.DATABASE_URL, pool_pre_ping=True, echo=True, echo_pool=True)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def get_session():
with Session(engine) as session:
yield session
def get_db() -> Generator:
try:
db = SessionLocal()
yield db
finally:
db.close()
I also tried using get_session instead of get_db, then the error is: no such table: users. Any ideas what is wrong with my setup?
if you have an already table in your database , you have to use MetaData to map them like:
from sqlalchemy import create_engine, MetaData
meta = MetaData(bind=engine)
meta.reflect(views=True)
table_name = meta.tables['table_name']
Session = sessionmaker(autocommit=False, autoflush=False, bind=engine)
session = Session()
def get_db():
db = session
try:
yield db
finally:
db.close()

Flask SQLAlchemy does not close MySQL database connections

I have a Flask app using Flask-SQLAlchemy with a MySQL database where the db is defined as the following:
db.py:
from flask_sqlalchemy import SQLAlchemy
db = SQLAlchemy()
main.py:
from db import db
app = Flask(__name__)
app.config['SQLALCHEMY_TRACK_MODIFICATIONS'] = False
app.config['SQLALCHEMY_DATABASE_URI'] = "mysql+pymysql://" + \
DB_USERNAME + ":" + DB_PASSWORD + "#" + DB_HOST + "/" + DB_DATABASE
db.init_app(app)
#app.teardown_appcontext
def teardown_db(error):
db.session.close()
db.engine.dispose()
user.py:
class User(db.Model):
id = db.Column(db.Integer, primary_key=True)
username = db.Column(db.String(80), unique=True, nullable=False)
email = db.Column(db.String(120), unique=True, nullable=False)
I query my database using models using either db.engine.execute() to write raw SQL queries where required or use the integrated Flask-SQLAlchemy APIs for reading data such as User.query.filter_by().all().
I write new data into the db using the following:
new_user_entry = User(username = "abc", email = "abc#example.com")
db.session.add(new_user_entry)
db.session.commit()
I am monitoring my MySQL server using show processlist and I notice that the database connections keep increasing by 2 for every single request that comes my way. The database connections seem to reset only when I stop the Flask process. With time, the MySQL server throws the below error:
`sqlalchemy.exc.TimeoutError: QueuePool limit of size 10 overflow 10 reached, connection timed out, timeout 30 (Background on this error at: http://sqlalche.me/e/3o7r)`
I am serving the app using gunicorn and gevent/eventlet with 2 worker processes. I use python3.
Am I missing something here? I tried ending the db session and disposing the engine, but this does not seem to work.
I finally found a fix to the above problem.
I used the declarative model defined in here instead of following the quickstart documentation for Flask-SQLAlchemy given here.
The changed files are as follows:
db.py:
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy.ext.declarative import declarative_base
engine = create_engine(DB_URI, convert_unicode=True)
db_session = scoped_session(sessionmaker(autocommit=False,
autoflush=False,
bind=engine))
Base = declarative_base()
Base.query = db_session.query_property()
def init_db():
import user
Base.metadata.create_all(bind=engine)
main.py:
from db import init_db, db_session
init_db()
#app.teardown_appcontext
def shutdown_session(exception=None):
db_session.remove()
user.py:
from sqlalchemy import Column, Integer, String
from data_models.db import Base
class User(Base):
id = db.Column(Integer, primary_key=True)
username = db.Column(String(80), unique=True, nullable=False)
email = db.Column(String(120), unique=True, nullable=False)
To query for records we could either use User.query.filter_by().all() or db_engine.execute().
To write new data into the database, we can use the following:
new_user_entry = User(username = "abc", email = "abc#example.com")
db_session.add(new_user_entry)
db_session.commit()
In case we need to close session before creating a new child process (what is recommended), this is what we should use:
db.session.remove()
db.engine.dispose()
Like
from multiprocessing import Process
from app import db
#app.route('/process/start/', methods = ["GET"])
def process_start():
db.session.remove()
db.engine.dispose()
p = Process(target = long_task)
p.start()
return 'started the long task'
def long_task():
'''
do long task
'''
Use with statement, there is a test:
def test():
with db.session() as dbss:
qss = models.WhateverModel.query.session
assert dbss == qss

How to enforce reads from read replica in Postgres using django and or flask?

Is enforcement of a reads from a read replica an application layer task?
i.e. I have a Postgres database and I have set up a read replica.
On the application side I have two connections one for the "write" database and one for the "read replica" database.
In my code if I do a "read" action I use the connection to the read replica. But when I go to insert or update I use the connection to the "write" database a.k.a. master.
Is there better with django or flask that this is managed automatically. i.e.
I would prefer to avoid specifying directly in code the connection to use and just have django or flask figure it out on their own.
Django
For this purpose django supports so called Database routers.
First create your custom router:
class CustomRouter:
def db_for_read(self, model, **hints):
return 'replica'
def db_for_write(self, model, **hints):
return 'master'
And configure django orm to use it like that.
DATABASES = {
'default': {},
'primary': {
'NAME': 'master',
'ENGINE': 'django.db.backends.mysql',
'USER': 'mysql_user',
'PASSWORD': 'spam',
},
'replica1': {
'NAME': 'replica',
'ENGINE': 'django.db.backends.mysql',
'USER': 'mysql_user',
'PASSWORD': 'eggs',
},
}
DATABASE_ROUTERS = ['path.to.CustomRouter']
The sample code was taken from the docs (it is worth reading!) and slightly adjusted.
SQLAlchemy (flask)
I went through the SQLAlchemy docs and found a link to this article, which describes how to implement djangos database router approach with SQLAlchemy.
You can use a custom session here to implement this properly.
The following snippets are picked from the linked article and are slightly adjusted.
Create your engines:
engines = {
'master': create_engine('postgresql://user:***#localhost:5432/master',
logging_name='master'),
'replica': create_engine('postgresql://user:***#localhost:5432/replica',
logging_name='replica'),
}
Create a custom session class:
class RoutingSession(Session):
def get_bind(self, mapper=None, clause=None):
if self._flushing:
return engines['master']
else:
return engines['replica']
And create your session like this:
Session = scoped_session(sessionmaker(class_=RoutingSession, autocommit=True))
Read the article for details and limitations.
here is concept solution for flask witch change
"""This is not the full code. We do a lot of stuff to clean up connections, particularly for unit testing."""
import sqlalchemy
from sqlalchemy.orm import Query, Session, scoped_session, sessionmaker
CONFIG_KEY_SQLALCHEMY_BINDS = 'SQLALCHEMY_BINDS'
CONFIG_KEY_SQLALCHEMY_RO_BINDS = 'SQLALCHEMY_READ_ONLY_BINDS'
class Config:
# These default values are for testing. In a deployed environment, they would be three separate instances.
SQLALCHEMY_DATABASE_URI = 'postgresql://localhost/branded_dev'
SQLALCHEMY_READ_ONLY_BINDS = {
'replica': 'postgresql://localhost/branded_dev',
'replica_analytics': 'postgresql://localhost/branded_dev'
}
class DBSessionFactory:
"""
A wrapper for getting db sessions from the primary and read replicas.
"""
def register(config):
self.engines = dict() # type: Dict[str, Engine]
self.read_only_engines = defaultdict(list) # type: Dict[str, Engine]
# The session factories to be used by scoped_session to connect
self.session_factories = dict() # Dict[str, sessionmaker]
# The scoped sessions for each connection.
self.scoped_sessions = dict() # Dict[str, scoped_session]
# The scoped sessions for each read only connection.
self.read_only_scoped_sessions = defaultdict(list) # Dict[str, List[scoped_session]]
# The primary connection
self.add_engine(
'primary', config.SQLALCHEMY_DATABASE_URI, config=config
)
# Other read-write dbs
for name, connect_url in config[CONFIG_KEY_SQLALCHEMY_BINDS].items():
self.add_engine(name, connect_url, config=config)
# Read replica binds
for name, connect_url in config[CONFIG_KEY_SQLALCHEMY_RO_BINDS].items():
self.add_engine(name, connect_url, config=config, read_only=True)
def add_engine(self, name: DBInstance, uri: str, config: Config, read_only=False) -> None:
"""Initialize a database connection and register it in the appropriate internal dicts."""
# Clean up existing engine if present
if self.engines.get(name) or self.read_only_engines.get(name):
self.session_factories[name].close_all()
engines = [self._create_engine(u, config) for u in uri] if isinstance(uri, list) \
else [self._create_engine(uri, config)]
for engine in engines:
self.session_factories[name] = sessionmaker(bind=engine, expire_on_commit=False)
scoped_session_instance = scoped_session(self.session_factories[name])
if read_only:
self.read_only_engines[name].append(engine)
self.read_only_scoped_sessions[name].append(scoped_session_instance)
else:
self.engines[name] = engine
self.scoped_sessions[name] = scoped_session_instance
def _create_engine(self, url: str, config: Config): # pylint: disable=no-self-use
"""wrapper to set up our connections"""
engine = sqlalchemy.create_engine(
url,
pool_size=config.SQLALCHEMY_POOL_SIZE,
pool_recycle=config.SQLALCHEMY_POOL_RECYCLE,
echo=config.SQLALCHEMY_ECHO,
pool_pre_ping=config.SQLALCHEMY_POOL_PRE_PING
)
#contextmanager
def session(self, engine: DBInstance=None) -> Generator[scoped_session, None, None]:
"""
Generate a session and yield it out.
After resuming, commit, unless an exception happens,
in which case we roll back.
:param engine: connection to use
:return: a generator for a scoped session
"""
session = self.raw_scoped_session(engine)
try:
yield session
session.commit()
except:
session.rollback()
raise
finally:
session.remove()
def read_only_session(self, engine: str=None) -> scoped_session:
"""
Return a session for a read-only db
:param engine: connection to use
:return: a Session via scoped_session
"""
if engine in self.read_only_engines:
return random.choice(self.read_only_scoped_sessions[engine])
else:
raise DBConfigurationException(
"Requested session for '{}', which is not bound in this app. Try: [{}]".
format(engine, ','.join(self.read_only_engines.keys()))
)
# The global db factory instance.
db = DBSessionFactory()
https://gist.github.com/jasonwalkeryung/5133383d66782461cdc3b4607ae35d98

How to rollback database changes during Flask testing

When running tests in pytest, the database is modified. What is the best way to undo changes to the database?
DBSession rollback
For those tests where I can access the backend directly, I currently use pytest fixture to start a new DBSession for every test function, and rollback the session at the end of it
#pytest.fixture(scope='session')
def db(app, request):
"""Session-wide test database."""
def teardown():
pass
_db = SQLAlchemy(app)
return _db
#pytest.fixture(scope='function')
def db_session(db, request):
"""Creates a new database session for a test."""
engine = create_engine(
TestConfig.SQLALCHEMY_DATABASE_URI,
connect_args={"options": "-c timezone=utc"})
DbSession = sessionmaker(bind=engine)
session = DbSession()
connection = engine.connect()
transaction = connection.begin()
options = dict(bind=connection, binds={})
session = db.create_scoped_session(options=options)
db.session = session
yield session
transaction.rollback()
connection.close()
session.remove()
In the test code, I simply use the fixture
def test_create_project(db_session):
project = _create_test_project(db_session)
assert project.project_id > 0
Flask / HTTP Testing
But for testing the API via Flask/HTTP, I cannot use db_session. Even when I create a fixture to explicitly DROP the test database and restore from production, it will not work because there is no direct database code
#pytest.fixture(scope='function')
def db_session_refresh(db, request):
"""Refresh the test database from production after running the test"""
engine = create_engine(
TestConfig.SQLALCHEMY_DATABASE_URI,
connect_args={"options": "-c timezone=utc"})
DbSession = sessionmaker(bind=engine)
session = DbSession()
connection = engine.connect()
transaction = connection.begin()
options = dict(bind=connection, binds={})
session = db.create_scoped_session(options=options)
db.session = session
yield session
transaction.rollback()
connection.close()
session.remove()
refresh_test_db_sql = """
SELECT pg_terminate_backend(pid)
FROM pg_stat_activity
WHERE datname = 'appdb_test';
DROP DATABASE appdb_test;
CREATE DATABASE appdb_test TEMPLATE appdb;
"""
engine.execute(refresh_test_db_sql)
Even if this works, it is inefficient to refresh the database for every function.
What is the proper/better way to run test that modifies the database?
as commented earlier - creation and destruction of the DB should be taken care out of you unit test and moved into a wrapper.
but to answer you question - I have the impression that you're using postgresql. try removing the database name from your connection string and connect to it directly and add it to the search_path

Categories