Tortoise.get_connection in pytest test raises a KeyError

Describe the bug My project uses some raw SQL. For this, I use Tortoise.get_connection() to get a DB connection. This fails when I run unit tests with pytest with a KeyError.

To Reproduce test_connection.py:

from tortoise import Tortoise
async def test_get_connection():
    c = Tortoise.get_connection('default')

$ pytest test_connection.py

Expected behavior The test should succeed without an exception. I also tried ‘models’ and ‘test’ for the connection name.

Additional context This seems to be caused by tortoise.contrib.test.initializer() which contains this code:

_CONNECTIONS = Tortoise._connections.copy()
_CONN_MAP = current_transaction_map.copy()
Tortoise.apps = {}
Tortoise._connections = {} 

[s3rius]

I’ve found a workaround. This is my create_db fixture.

import pytest
from tortoise.contrib.test import finalizer, initializer
from tortoise.transactions import current_transaction_map


@pytest.fixture(autouse=True)
@pytest.mark.asyncio
def create_db() -> Generator[None, None, None]:
    """
    Fixture that initializes database.

    :yield: Nothing.
    """
    db_url = os.environ.get("PROJECT_TEST_DB", "sqlite://:memory:")
    initializer(MODELS_MODULES, db_url=db_url, app_label="models")
    current_transaction_map["default"] = current_transaction_map["models"]

    yield

    finalizer()

And the function looks like this:

from tortoise.transactions import atomic

@atomic("default")
async def grow(message: Message) -> None:
  ... # code

[DrJackilD]

@janpascal I have the same problems with pytest (and to be honest I don’t want to switch to other test framework or unittest class-based style), but I found a solution to use a connection name in transaction.atomic decorator. The general difference between my and your setup is that I use module names during Tortoise initialize (in conftest and the application)

Like this:

initializer(["app.db.models"], db_url=str(settings.DB_URL), app_label="models", loop=event_loop)

In that case Tortoise will create current_transaction_map[name] link in tortoise.Tortoise._init_connections and you can get it with @transaction.atomic("models").

Hope this will help.

UPD: No, it’s not, now I’ve got a problem in the application 😅

UPD2 I made this ugly hack to get what I want. I’m using TESTING in my settings to understand is this a test environment, so I’ve added this to the decorator:

@transaction.atomic("models" if settings.TESTING else None)

Looks awful, but it works.

I believe the reason that it works is here:

# tortoise.Tortoise._init_connections
...
await connection.create_connection(with_db=True)
cls._connections[name] = connection
current_transaction_map[name] = ContextVar(name, default=connection)

The we clear Tortoise._connections in tortoise.contrib.test.initializer, but current_transaction_map left untouched.