embedchain

1
import os
2

3
import pytest
4
from sqlalchemy import MetaData, create_engine
5
from sqlalchemy.orm import sessionmaker
6

7

8
@pytest.fixture(autouse=True)
9
def clean_db():
10
    db_path = os.path.expanduser("~/.embedchain/embedchain.db")
11
    db_url = f"sqlite:///{db_path}"
12
    engine = create_engine(db_url)
13
    metadata = MetaData()
14
    metadata.reflect(bind=engine)  # Reflect schema from the engine
15
    Session = sessionmaker(bind=engine)
16
    session = Session()
17

18
    try:
19
        # Iterate over all tables in reversed order to respect foreign keys
20
        for table in reversed(metadata.sorted_tables):
21
            if table.name != "alembic_version":  # Skip the Alembic version table
22
                session.execute(table.delete())
23
        session.commit()
24
    except Exception as e:
25
        session.rollback()
26
        print(f"Error cleaning database: {e}")
27
    finally:
28
        session.close()
29

30

31
@pytest.fixture(autouse=True)
32
def disable_telemetry():
33
    os.environ["EC_TELEMETRY"] = "false"
34
    yield
35
    del os.environ["EC_TELEMETRY"]
36