wehub-resource-sync
40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
from collections.abc import Iterator
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from _data import SEED_SPANS_PER_TRACE, SEED_TRACES, seed_traces
|
|
|
|
import mlflow
|
|
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def bench_dir(tmp_path_factory: pytest.TempPathFactory) -> Path:
|
|
return tmp_path_factory.mktemp("bench")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def store(bench_dir: Path) -> SqlAlchemyStore:
|
|
db_uri = f"sqlite:///{bench_dir / 'mlflow.db'}"
|
|
(bench_dir / "artifacts").mkdir()
|
|
artifact_root = (bench_dir / "artifacts").as_uri()
|
|
return SqlAlchemyStore(db_uri, artifact_root)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def experiment_id(store: SqlAlchemyStore) -> str:
|
|
return str(store.create_experiment("bench"))
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def seeded(store: SqlAlchemyStore, experiment_id: str) -> list[str]:
|
|
return seed_traces(store, experiment_id, SEED_TRACES, SEED_SPANS_PER_TRACE)
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def e2e_setup(bench_dir: Path) -> Iterator[None]:
|
|
mlflow.set_tracking_uri(f"sqlite:///{bench_dir / 'e2e.db'}")
|
|
mlflow.set_experiment("bench_e2e")
|
|
yield
|
|
mlflow.flush_trace_async_logging(terminate=True)
|