How to use testcontainers to test sqlalchemy on real database

in previous article, I described how to use sqlalchemy and how to do unit test with python build in db sqlite, sometimes with legacy code, we have to execute raw sql with sqlalchemy engine and some of the raw sql is using database specific functions, to test these logic, we could not spin up sqlite , we have to launch real database. with testcontainers-python, things get really easy.

below is the code example from testcontainers-python's webiste to demo how to use it

import sqlalchemy
from testcontainers.mysql import MySqlContainer

with MySqlContainer('mysql:5.7.17') as mysql:
    engine = sqlalchemy.create_engine(mysql.get_connection_url())
    version, = engine.execute("select version()").fetchone()
    print(version)  # 5.7.17

in my case, I'm creating engine based on the DB_URL system env, so my first attempt is like below

@pytest.fixture(scope="function", autouse=True)
def init_db():
    from testcontainers.mysql import MySqlContainer

    with MySqlContainer("mysql:5.7.24") as mysql:
        os.environ["DB_URL"] = mysql.get_connection_url()
        Base.metadata.create_all(engine_pair)
        yield
        os.environ["DB_URL"] = "sqlite:///:memory:pair"

when try the unit test, it failed, and reported it's still using sqlite db rather than mysql. then I figure out my sqlalchemy engine is created in the __init__.py statically, if I want to override it, I must reload the module after system env changed, importlib could help in this context. I need to define the system env DB_URL first, then reload the module to re-create my sqlalchemy engine, after test complete, I could reset system env DB_URL and reload the module again to clean up

@pytest.fixture(scope="function", autouse=True)
def init_db():
    from testcontainers.mysql import MySqlContainer

    with MySqlContainer("mysql:5.7.24") as mysql:
        import os
        import importlib
        import db
        from db import database
        from db.models import acme_record
        os.environ["DB_URL"] = mysql.get_connection_url()
        importlib.reload(db)
        importlib.reload(database)
        importlib.reload(acme_record)
        from db import engine
        from db.models import Base
        Base.metadata.create_all(engine)
        yield
    # set env back
    import os
    import importlib
    import db
    from db import database

    os.environ["DB_URL"] = "sqlite:///:memory"
    importlib.reload(db)
    importlib.reload(database)