wehub-resource-sync
28 lines
678 B
Python
28 lines
678 B
Python
import pytest
|
|
|
|
import mlflow
|
|
|
|
from tests.tracing.helper import purge_traces
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_mlflow_autolog_and_traces():
|
|
yield
|
|
mlflow.pydantic_ai.autolog(disable=True)
|
|
purge_traces()
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clear_autolog_state():
|
|
from mlflow.utils.autologging_utils import AUTOLOGGING_INTEGRATIONS
|
|
|
|
for key in AUTOLOGGING_INTEGRATIONS.keys():
|
|
AUTOLOGGING_INTEGRATIONS[key].clear()
|
|
mlflow.utils.import_hooks._post_import_hooks = {}
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def mock_creds(monkeypatch):
|
|
monkeypatch.setenv("OPENAI_API_KEY", "my-secret-key")
|
|
monkeypatch.setenv("GEMINI_API_KEY", "my-secret-key")
|