1  import pytest
 2  
 3  import mlflow
 4  
 5  from tests.tracing.helper import purge_traces
 6  
 7  
 8  @pytest.fixture(autouse=True)
 9  def reset_mlflow_autolog_and_traces():
10      yield
11      mlflow.pydantic_ai.autolog(disable=True)
12      purge_traces()
13  
14  
15  @pytest.fixture(autouse=True)
16  def clear_autolog_state():
17      from mlflow.utils.autologging_utils import AUTOLOGGING_INTEGRATIONS
18  
19      for key in AUTOLOGGING_INTEGRATIONS.keys():
20          AUTOLOGGING_INTEGRATIONS[key].clear()
21      mlflow.utils.import_hooks._post_import_hooks = {}
22  
23  
24  @pytest.fixture(autouse=True)
25  def mock_creds(monkeypatch):
26      monkeypatch.setenv("OPENAI_API_KEY", "my-secret-key")
27      monkeypatch.setenv("GEMINI_API_KEY", "my-secret-key")