1  import os
 2  from datetime import timedelta
 3  
 4  import pytest
 5  from testcontainers.compose import DockerCompose
 6  from testcontainers.core.wait_strategies import HttpWaitStrategy
 7  
 8  import mlflow
 9  
10  
11  @pytest.mark.parametrize(
12      "compose_file",
13      [
14          "docker-compose.mssql-test.yaml",
15          "docker-compose.mysql-test.yaml",
16          "docker-compose.postgres-test.yaml",
17      ],
18  )
19  def test_backend_and_artifact_store_integration(compose_file):
20      compose = DockerCompose(
21          context=os.path.dirname(os.path.abspath(__file__)),
22          compose_file_name=[compose_file],
23      )
24      # Configure wait strategy before starting containers
25      compose.waiting_for({
26          "mlflow": HttpWaitStrategy(5000, "/health")
27          .for_status_code(200)
28          .with_startup_timeout(timedelta(minutes=5))
29      })
30  
31      with compose:
32          base_url = "http://localhost:5000"
33  
34          mlflow.set_tracking_uri(base_url)
35          mlflow.set_experiment("integration-test")
36  
37          @mlflow.trace
38          def predict(model_input: list[str]) -> list[str]:
39              return model_input
40  
41          with mlflow.start_run():
42              mlflow.log_param("param", 1)
43              mlflow.log_metric("metric", 1.0)
44              mlflow.pyfunc.log_model(
45                  name="test_model",
46                  python_model=predict,
47                  input_example=["a", "b", "c"],
48              )