Function mock_db_session

tests/main/sql_chat/test_sql_chat_agent.py:55–84

()

@pytest.fixture
def mock_db_session() -> Session:
    # Create an in-memory SQLite database
    engine = create_engine("sqlite:///:memory:", echo=False)
    Base.metadata.create_all(engine)

    Session = sessionmaker(bind=engine)
    session = Session()

    # Insert data
    sales_dept = Department(id=1, name="Sales")
    marketing_dept = Department(id=2, name="Marketing")

    alice = Employee(id=1, name="Alice", department=sales_dept)
    bob = Employee(id=2, name="Bob", department=marketing_dept)

    sale1 = Sale(id=1, amount=100, employee=alice)
    sale2 = Sale(id=2, amount=500, employee=bob)

    session.add(sales_dept)
    session.add(marketing_dept)
    session.add(alice)
    session.add(bob)
    session.add(sale1)
    session.add(sale2)

    session.commit()

    yield session  # this is where the fixture's value comes from

    session.close()


@pytest.fixture