Change staging test to new alias based test

burakince · Dec 13, 2024 · 2ce5f5c · 2ce5f5c
1 parent 36cbdf9
commit 2ce5f5c
Show file tree

Hide file tree

Showing 10 changed files with 57 additions and 49 deletions.
diff --git a/.dockerignore b/.dockerignore
@@ -3,6 +3,7 @@
 .git
 .github
 .pytest_cache
+.vscode
 
 # Dot Files
 .pre-commit-config.yaml

diff --git a/.env b/.env
@@ -2,3 +2,5 @@ POSTGRES_VERSION=15.10
 MYSQL_VERSION=8.4.3
 MINIO_VERSION=RELEASE.2024-11-07T00-52-20Z
 FAKE_GCS_SERVER_VERSION=1.50.2
+AZURITE_VERSION=3.33.0
+AZURE_STORAGE_BLOB_VERSION=12.26.0
diff --git a/docker-compose.azure-mysql-test.yaml b/docker-compose.azure-mysql-test.yaml
@@ -7,6 +7,9 @@ services:
     build:
       context: test-containers/azurite
       dockerfile: Dockerfile
+      args:
+        AZURITE_VERSION: ${AZURITE_VERSION}
+        AZURE_STORAGE_BLOB_VERSION: ${AZURE_STORAGE_BLOB_VERSION}
     command: sh -c "sh /run.sh | azurite -l /data --blobHost 0.0.0.0 --queueHost 0.0.0.0 --loose --skipApiVersionCheck"
     environment:
       AZURE_STORAGE_ACCOUNT: devstoreaccount1

diff --git a/docker-compose.azure-postgres-test.yaml b/docker-compose.azure-postgres-test.yaml
@@ -7,6 +7,9 @@ services:
     build:
       context: test-containers/azurite
       dockerfile: Dockerfile
+      args:
+        AZURITE_VERSION: ${AZURITE_VERSION}
+        AZURE_STORAGE_BLOB_VERSION: ${AZURE_STORAGE_BLOB_VERSION}
     command: sh -c "sh /run.sh | azurite -l /data --blobHost 0.0.0.0 --queueHost 0.0.0.0 --loose --skipApiVersionCheck"
     environment:
       AZURE_STORAGE_ACCOUNT: devstoreaccount1

diff --git a/tests/test_aws_cloud_mysql_integration.py b/tests/test_aws_cloud_mysql_integration.py
@@ -13,7 +13,7 @@ def test_mysql_backended_aws_simulation_model_upload_and_access_via_api(
     with DockerCompose(
         context=".",
         compose_file_name=["docker-compose.aws-mysql-test.yaml"],
-        pull=True,
+        # pull=True,
         build=True,
     ) as compose:
         mlflow_host = compose.get_service_host("mlflow", 8080)
@@ -43,16 +43,15 @@ def test_mysql_backended_aws_simulation_model_upload_and_access_via_api(
             model_details = mlflow.register_model(model_uri, model_name)
 
             client = MlflowClient()
-            client.transition_model_version_stage(
+            client.set_registered_model_alias(
                 name=model_details.name,
+                alias=stage_name,
                 version=model_details.version,
-                stage=stage_name,
             )
 
-        params = {"name": model_name, "stages": stage_name}
-        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/get-latest-versions"
-        r = requests.get(url=latest_version_url, params=params)
+        params = {"name": model_name, "alias": stage_name}
+        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/alias"
+        r = requests.get(url=latest_version_url, params=params, timeout=300)
 
-        print(r.json())
-        assert "1" == r.json()["model_versions"][0]["version"]
-        assert "READY" == r.json()["model_versions"][0]["status"]
+        assert "1" == r.json()["model_version"]["version"]
+        assert "READY" == r.json()["model_version"]["status"]
diff --git a/tests/test_aws_cloud_postgres_integration.py b/tests/test_aws_cloud_postgres_integration.py
@@ -13,7 +13,7 @@ def test_postgres_backended_aws_simulation_model_upload_and_access_via_api(
     with DockerCompose(
         context=".",
         compose_file_name=["docker-compose.aws-postgres-test.yaml"],
-        pull=True,
+        # pull=True,
         build=True,
     ) as compose:
         mlflow_host = compose.get_service_host("mlflow", 8080)
@@ -43,15 +43,15 @@ def test_postgres_backended_aws_simulation_model_upload_and_access_via_api(
             model_details = mlflow.register_model(model_uri, model_name)
 
             client = MlflowClient()
-            client.transition_model_version_stage(
+            client.set_registered_model_alias(
                 name=model_details.name,
+                alias=stage_name,
                 version=model_details.version,
-                stage=stage_name,
             )
 
-        params = {"name": model_name, "stages": stage_name}
-        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/get-latest-versions"
-        r = requests.get(url=latest_version_url, params=params)
+        params = {"name": model_name, "alias": stage_name}
+        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/alias"
+        r = requests.get(url=latest_version_url, params=params, timeout=300)
 
-        assert "1" == r.json()["model_versions"][0]["version"]
-        assert "READY" == r.json()["model_versions"][0]["status"]
+        assert "1" == r.json()["model_version"]["version"]
+        assert "READY" == r.json()["model_version"]["status"]
diff --git a/tests/test_azure_cloud_mysql_integration.py b/tests/test_azure_cloud_mysql_integration.py
@@ -14,7 +14,7 @@ def test_mysql_backended_azure_simulation_model_upload_and_access_via_api(
     with DockerCompose(
         context=".",
         compose_file_name=["docker-compose.azure-mysql-test.yaml"],
-        pull=True,
+        # pull=True,
         build=True,
     ) as compose:
         mlflow_host = compose.get_service_host("mlflow", 8080)
@@ -54,18 +54,18 @@ def test_mysql_backended_azure_simulation_model_upload_and_access_via_api(
             model_details = mlflow.register_model(model_uri, model_name)
 
             client = MlflowClient()
-            client.transition_model_version_stage(
+            client.set_registered_model_alias(
                 name=model_details.name,
+                alias=stage_name,
                 version=model_details.version,
-                stage=stage_name,
             )
 
-        params = {"name": model_name, "stages": stage_name}
-        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/get-latest-versions"
-        r = requests.get(url=latest_version_url, params=params)
+        params = {"name": model_name, "alias": stage_name}
+        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/alias"
+        r = requests.get(url=latest_version_url, params=params, timeout=300)
 
-        assert "1" == r.json()["model_versions"][0]["version"]
-        assert "READY" == r.json()["model_versions"][0]["status"]
+        assert "1" == r.json()["model_version"]["version"]
+        assert "READY" == r.json()["model_version"]["status"]
 
 
 def connection_str(host, blob_port, queue_port):

diff --git a/tests/test_azure_cloud_postgres_integration.py b/tests/test_azure_cloud_postgres_integration.py
@@ -14,7 +14,7 @@ def test_postgres_backended_azure_simulation_model_upload_and_access_via_api(
     with DockerCompose(
         context=".",
         compose_file_name=["docker-compose.azure-postgres-test.yaml"],
-        pull=True,
+        # pull=True,
         build=True,
     ) as compose:
         mlflow_host = compose.get_service_host("mlflow", 8080)
@@ -58,18 +58,18 @@ def test_postgres_backended_azure_simulation_model_upload_and_access_via_api(
             model_details = mlflow.register_model(model_uri, model_name)
 
             client = MlflowClient()
-            client.transition_model_version_stage(
+            client.set_registered_model_alias(
                 name=model_details.name,
+                alias=stage_name,
                 version=model_details.version,
-                stage=stage_name,
             )
 
-        params = {"name": model_name, "stages": stage_name}
-        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/get-latest-versions"
-        r = requests.get(url=latest_version_url, params=params)
+        params = {"name": model_name, "alias": stage_name}
+        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/alias"
+        r = requests.get(url=latest_version_url, params=params, timeout=300)
 
-        assert "1" == r.json()["model_versions"][0]["version"]
-        assert "READY" == r.json()["model_versions"][0]["status"]
+        assert "1" == r.json()["model_version"]["version"]
+        assert "READY" == r.json()["model_version"]["status"]
 
 
 def connection_str(host, blob_port, queue_port):

diff --git a/tests/test_gcp_cloud_mysql_integration.py b/tests/test_gcp_cloud_mysql_integration.py
@@ -13,7 +13,7 @@ def test_mysql_backended_gcp_simulation_model_upload_and_access_via_api(
     with DockerCompose(
         context=".",
         compose_file_name=["docker-compose.gcp-mysql-test.yaml"],
-        pull=True,
+        # pull=True,
         build=True,
     ) as compose:
         mlflow_host = compose.get_service_host("mlflow", 8080)
@@ -42,15 +42,15 @@ def test_mysql_backended_gcp_simulation_model_upload_and_access_via_api(
             model_details = mlflow.register_model(model_uri, model_name)
 
             client = MlflowClient()
-            client.transition_model_version_stage(
+            client.set_registered_model_alias(
                 name=model_details.name,
+                alias=stage_name,
                 version=model_details.version,
-                stage=stage_name,
             )
 
-        params = {"name": model_name, "stages": stage_name}
-        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/get-latest-versions"
-        r = requests.get(url=latest_version_url, params=params)
+        params = {"name": model_name, "alias": stage_name}
+        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/alias"
+        r = requests.get(url=latest_version_url, params=params, timeout=300)
 
-        assert "1" == r.json()["model_versions"][0]["version"]
-        assert "READY" == r.json()["model_versions"][0]["status"]
+        assert "1" == r.json()["model_version"]["version"]
+        assert "READY" == r.json()["model_version"]["status"]
diff --git a/tests/test_gcp_cloud_postgres_integration.py b/tests/test_gcp_cloud_postgres_integration.py
@@ -13,7 +13,7 @@ def test_postgres_backended_gcp_simulation_model_upload_and_access_via_api(
     with DockerCompose(
         context=".",
         compose_file_name=["docker-compose.gcp-postgres-test.yaml"],
-        pull=True,
+        # pull=True,
         build=True,
     ) as compose:
         mlflow_host = compose.get_service_host("mlflow", 8080)
@@ -42,15 +42,15 @@ def test_postgres_backended_gcp_simulation_model_upload_and_access_via_api(
             model_details = mlflow.register_model(model_uri, model_name)
 
             client = MlflowClient()
-            client.transition_model_version_stage(
+            client.set_registered_model_alias(
                 name=model_details.name,
+                alias=stage_name,
                 version=model_details.version,
-                stage=stage_name,
             )
 
-        params = {"name": model_name, "stages": stage_name}
-        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/get-latest-versions"
-        r = requests.get(url=latest_version_url, params=params)
+        params = {"name": model_name, "alias": stage_name}
+        latest_version_url = f"{base_url}/api/2.0/mlflow/registered-models/alias"
+        r = requests.get(url=latest_version_url, params=params, timeout=300)
 
-        assert "1" == r.json()["model_versions"][0]["version"]
-        assert "READY" == r.json()["model_versions"][0]["status"]
+        assert "1" == r.json()["model_version"]["version"]
+        assert "READY" == r.json()["model_version"]["status"]
-Original file line number
+Diff line change
@@ Expand Up / @@ -3,6 +3,7 @@ @@
     .git
     .github
     .pytest_cache
+    .vscode
     # Dot Files
     .pre-commit-config.yaml
@@ Expand Down @@