Search code examples
pythonmysqlsqlalchemypython-asynciofastapi

Running Pytest test cases in transaction isolation in a FastAPI setup


I have a FastAPI application, with MySQL and asyncio.

I have been trying to integrate some test cases with my application, with the ability to rollback the changes after every test case, so that all test cases can run in isolation.

This is how my controller is set up, with a DB dependency getting injected.

from sqlalchemy.ext.asyncio import create_async_engine

async def get_db_connection_dependency():
    engine = create_async_engine("mysql+aiomysql://root:root@mysql8:3306/user_db")
    connection = engine.connect()
    return connection

class UserController:
   async def create_user(
            self,
            request: Request,
            connection: AsyncConnection = Depends(get_db_connection_dependency)

    ) -> JSONResponse:
        
        # START TRANSACTION
        await connection.__aenter__()
        transaction = connection.begin()
        await transaction.__aenter__()

        try:
            do_stuff()
        except:
            await transaction.rollback()
        else:
            await transaction.commit()
        finally:
            await connection.close()
      
        # END TRANSACTION
        
        return JSONResponse(status_code=201)

I have a test case written using Pytest like so

import pytest

app = FastAPI()

@pytest.fixture()
def client():
    with TestClient(app=app) as c:
        yield c

class TestUserCreation:
    CREATE_USER_URL = "/users/create"
    
    def test_create_user(self, client):
        response = client.post(self.CREATE_USER_URL, json={"name": "John"})
        assert response.status_code == 201

This test case works and persists the newly created user in the DB, but like I said earlier, I want to rollback the changes automatically once the test case finishes.

I have checked a few resources online, but none of them were helpful.

  1. This link talks about using factory objects, but I can't use factory objects here because my controller requires the DB connection as a dependency. Plus, the controller itself is updating the DB, and not a "mocked" factory object.

  2. I then searched for ways to inject the dependency manually. This was in the hopes that if I can create a connection manually BEFORE calling the API in my test case and inject it as the required dependency, then I can also forcefully rollback the transaction AFTER the API finishes.

    • So, I came across this, which talks about a way to get a dependency to use outside of a controller, but not how to inject it into the controller manually.
  3. The official FastAPI docs aren't very exhaustive on how to rollback persisted data in a DB-related test case.

The only way I can think of is to not inject the DB connection as a dependency into the controller, but attach it to the Starlette request object in the request middleware. And then in the response middleware, depending on an env var (test vs prod), I can ALWAYS rollback if the var is test.

But this seems over-engineering to me for a very fundamental requirement of a robust testing suite.

Is there any readily-available, built-in way to do this in FastAPI? Or is there any other library or package available that can do it for me?

If Pytest isn't the best suited framework for this, I'm more than happy to change it to something more suitable.

Appreciate any help I can get. Thank you!


Solution

  • I have solved this.

    I moved all transaction and connection handling into the dependency, with the controller doing nothing with the connection object. This meant that my actual API dependency can have the usual commit and rollback functions called for their appropriate scenarios, and my override dependency in my test case will always call rollback no matter what.

    My main problem was that I was letting my controller do some of the DB connection stuff which complicated things.