FastAPI psycopg2 error when creating database for tests - python

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()

Related

How to use dependency injection with simple functions in Fast API?

I have a simple function in Fast API app that gets called by a cron, something like this (note that this is not a get or post Fast API method but a simple function):
def read_user(user_id: int, db: Session = Depends(get_db)):
db_user = crud.get_user(db, user_id=user_id)
where get_db is:
from database import SessionLocal
def get_db():
db = SessionLocal()
try:
yield db
finally:
db.close()
And database.py is:
from sqlalchemy import create_engine
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
SQLALCHEMY_DATABASE_URL = "://url-to-db"
engine = create_engine(
SQLALCHEMY_DATABASE_URL
)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
Base = declarative_base()
Since crud method takes db (object of SessionLocal yielded by get_db) as a param, how do I use it as a dependency injection with my function just like we use in GET or POST as shown in code read_user method above.

Can we use a different DB for running pytest?

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.

How to verify session and engine objects are able to connect

The lines below declare the engine and session objects.
from sqlalchemy import create_engine
from sqlalchemy.orm import scoped_session, sessionmaker
URL = 'mysql://root:pass#12.34.56.78/mydb'
engine = create_engine(URL, echo=False)
session = scoped_session(sessionmaker())
session.remove()
session.configure(bind=engine, autoflush=False, expire_on_commit=False)
I would like to know how engine and session objects could be used to check if they are able to connect to the database with URL, database name and credentials provided.

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()
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.

Declarative SQLAlchemy CREATE SQLITE in memory tables

This is how I setup my database for an application (in Flask):
from sqlalchemy.engine import create_engine
from sqlalchemy.orm import scoped_session, create_session
from sqlalchemy.ext.declarative import declarative_base
engine = None
db_session = scoped_session(lambda: create_session(bind=engine,
autoflush=False, autocommit=False, expire_on_commit=True))
Base = declarative_base()
Base.query = db_session.query_property()
def init_engine(uri, **kwargs):
global engine
engine = create_engine(uri, **kwargs)
Base.metadata.create_all(bind=engine)
return engine
If I connect to a file database that has had tables created already, everything works fine, but using sqlite:///:memory: as a target database gives me:
OperationalError: (OperationalError) no such table: users u'DELETE FROM users' ()
when querying like so for ex.:
UsersTable.query.delete()
db_session.commit()
I am accessing this code from a unit test. What is the problem?
Thanks
Edit:
Working setup of the application:
app = Flask(__name__)
app.config.from_object(__name__)
app.secret_key = 'XXX'
# presenters
from presenters.users import users
# register modules (presenters)
app.register_module(users)
# initialize the database
init_engine(db)
The code you posted doesn't contain any table/class declaration. Are you sure that the declaration is done before init_engine() is called?

Categories