CODEWITHSUNDEEP
pythonunit-testingpytestpsycopg2
Dmytro Chasovskyi
Dmytro Chasovskyi

Reputation: 3651

What is the correct way to mock psycopg2 with pytest?

I need to simulate DB connection without actual connection. All answers I found are trying to mock methods in different ways, connect to docker db, connect to actual PostgreSQL running locally. I believe I need mocking variant but I cannot formulate in my head how should I mock. Am I missing something? Am I moving into wrong direction?

I use PostgreSQL and psycopg2. Package psycopg2-binary

Database connection:

import os

import psycopg2
from loguru import logger
from psycopg2.extensions import parse_dsn


def init_currency_history_table(cursor):
    create_users_table_query = """
        CREATE TABLE IF NOT EXISTS history(
          id BIGINT PRIMARY KEY NOT NULL,
          event TEXT,
          creation_date TIMESTAMPTZ DEFAULT NOW()
        );
    """
    cursor.execute(create_users_table_query)


def load_db(db_url):
    db = psycopg2.connect(**db_url)
    db.autocommit = True
    return db


class PostgresqlApi(object):

    def __init__(self, load=load_db):
        logger.info(os.environ.get('DATABASE_URL'))
        db_url = parse_dsn(os.environ.get('DATABASE_URL'))
        db_url['sslmode'] = 'require'
        logger.info('HOST: {0}'.format(db_url.get('host')))
        self.db = load_db(db_url)
        self.cursor = self.db.cursor()

        init_currency_history_table(self.cursor)
        self.db.commit()

    def add_event(self, *, event):
        insert_event_table = """
            INSERT INTO history (event) VALUES (%s);
        """
        self.cursor.execute(insert_event_table, (event))

    def events(self):
        select_event_table = """SELECT * FROM event;"""
        self.cursor.execute(select_event_table)
        return self.cursor.fetchall()

    def close(self):
        self.cursor.close()
        self.db.close()

I use DB for Falcon API.

from fastapi import Depends, FastAPI, HTTPException, status
from fastapi.security import HTTPBasic, HTTPBasicCredentials
from decimal import Decimal, getcontext

from db import PostgresqlApi

app = FastAPI()
security = HTTPBasic()
database = None


def db_connection():
    global database
    if not database:
        database = PostgresqlApi()
    return database

def check_basic_auth_creds(credentials: HTTPBasicCredentials = Depends(security)):
    correct_username = secrets.compare_digest(credentials.username, os.environ.get('APP_USERNAME'))
    correct_password = secrets.compare_digest(credentials.password, os.environ.get('APP_PASSWORD'))
    if not (correct_username and correct_password):
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username and password",
            headers={'WWW-Authenticate': 'Basic'}
        )
    return credentials

@app.get("/currencies")
def read_currencies(credentials: HTTPBasicCredentials = Depends(check_basic_auth_creds)):
    db = db_connection()
    return {'get events': 'ok'}

I have tried different methods and plugins. Among others arepytest-pgsql, pytest-postgresql.

Upvotes: 5

Views: 3812

Answers (1)

Dmytro Chasovskyi
Dmytro Chasovskyi

Reputation: 3651

The solution I landed at is below.

  • Created fake class that has exactly structure of PostgresqlApi. (see implementation below)
  • Created fixture for db_connection method. (see implementation below)

Fake class implementation

class FakePostgresqlApi(PostgresqlApi):

    event_list = []

    def __init__(self):
        pass

    def add_event(self, *, event):
        self.event_list.append([1, 'magic trick', 1653630607])

    def events(self):
        return self.event_list

    def close(self):
        self.event_list.clear()

Fixture

from unittest.mock import MagicMock

@pytest.fixture
def mock_db_connection(mocker):
    mocker.patch('src.main.db_connection', MagicMock(return_value=FakePostgresqlApi()))

The test itself was:

def test_read_events(mock_db_connection):
   # Do whatever I need here, in my case call Falcon API test client

Upvotes: 3

Related Questions