wehub-resource-sync
49 lines
1.3 KiB
Python
49 lines
1.3 KiB
Python
import os
|
|
from datetime import timedelta
|
|
|
|
import pytest
|
|
from testcontainers.compose import DockerCompose
|
|
from testcontainers.core.wait_strategies import HttpWaitStrategy
|
|
|
|
import mlflow
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"compose_file",
|
|
[
|
|
"docker-compose.mssql-test.yaml",
|
|
"docker-compose.mysql-test.yaml",
|
|
"docker-compose.postgres-test.yaml",
|
|
],
|
|
)
|
|
def test_backend_and_artifact_store_integration(compose_file):
|
|
compose = DockerCompose(
|
|
context=os.path.dirname(os.path.abspath(__file__)),
|
|
compose_file_name=[compose_file],
|
|
)
|
|
# Configure wait strategy before starting containers
|
|
compose.waiting_for({
|
|
"mlflow": HttpWaitStrategy(5000, "/health")
|
|
.for_status_code(200)
|
|
.with_startup_timeout(timedelta(minutes=5))
|
|
})
|
|
|
|
with compose:
|
|
base_url = "http://localhost:5000"
|
|
|
|
mlflow.set_tracking_uri(base_url)
|
|
mlflow.set_experiment("integration-test")
|
|
|
|
@mlflow.trace
|
|
def predict(model_input: list[str]) -> list[str]:
|
|
return model_input
|
|
|
|
with mlflow.start_run():
|
|
mlflow.log_param("param", 1)
|
|
mlflow.log_metric("metric", 1.0)
|
|
mlflow.pyfunc.log_model(
|
|
name="test_model",
|
|
python_model=predict,
|
|
input_example=["a", "b", "c"],
|
|
)
|