Add tox for running MLflow UATs

canonical · Sep 26, 2024 · acda82c · acda82c
1 parent 67b6f47
commit acda82c
Showing 1 changed file with 28 additions and 8 deletions.
diff --git a/tests/integration/test_bundle_deployment.py b/tests/integration/test_bundle_deployment.py
@@ -1,22 +1,30 @@
 import os
+
 import aiohttp
 import lightkube
 import pytest
-from pytest_operator.plugin import OpsTest
 from lightkube.resources.core_v1 import Service
+from pytest_operator.plugin import OpsTest
 
 # Environment variables
-KUBEFLOW_CHANNEL = os.environ.get("KUBEFLOW_CHANNEL", "1.9/stable")  # Default to '1.9/stable' if not set
-RESOURCE_DISPATCHER_CHANNEL = os.environ.get("RESOURCE_DISPATCHER_CHANNEL", "2.0/stable")  # Default to '2.0/stable' if not set
+KUBEFLOW_CHANNEL = os.environ.get(
+    "KUBEFLOW_CHANNEL", "1.9/stable"
+)  # Default to '1.9/stable' if not set
+RESOURCE_DISPATCHER_CHANNEL = os.environ.get(
+    "RESOURCE_DISPATCHER_CHANNEL", "2.0/stable"
+)  # Default to '2.0/stable' if not set
+
 
 @pytest.fixture()
 def lightkube_client() -> lightkube.Client:
     client = lightkube.Client(field_manager="kubeflow")
     return client
 
+
 @pytest.fixture
 def bundle_path() -> str:
-    return os.environ.get("BUNDLE_PATH").replace("\"", "")
+    return os.environ.get("BUNDLE_PATH").replace('"', "")
+
 
 async def deploy_bundle(ops_test: OpsTest, bundle_path, trust: bool) -> None:
     """Deploy a bundle from file using juju CLI."""
@@ -27,6 +35,7 @@ async def deploy_bundle(ops_test: OpsTest, bundle_path, trust: bool) -> None:
     print(stdout)
     assert retcode == 0, f"Deploy failed: {(stderr or stdout).strip()}"
 
+
 class TestCharm:
     @pytest.mark.abort_on_fail
     async def test_bundle_deployment_works(self, ops_test: OpsTest, lightkube_client, bundle_path):
@@ -49,9 +58,15 @@ async def test_bundle_deployment_works(self, ops_test: OpsTest, lightkube_client
 
         # Relate services as per Juju integrations
         await ops_test.model.relate("mlflow-server:secrets", "resource-dispatcher:secrets")
-        await ops_test.model.relate("mlflow-server:pod-defaults", "resource-dispatcher:pod-defaults")
-        await ops_test.model.relate("mlflow-minio:object-storage", "kserve-controller:object-storage")
-        await ops_test.model.relate("kserve-controller:service-accounts", "resource-dispatcher:service-accounts")
+        await ops_test.model.relate(
+            "mlflow-server:pod-defaults", "resource-dispatcher:pod-defaults"
+        )
+        await ops_test.model.relate(
+            "mlflow-minio:object-storage", "kserve-controller:object-storage"
+        )
+        await ops_test.model.relate(
+            "kserve-controller:service-accounts", "resource-dispatcher:service-accounts"
+        )
         await ops_test.model.relate("kserve-controller:secrets", "resource-dispatcher:secrets")
         await ops_test.model.relate("mlflow-server:ingress", "istio-pilot:ingress")
         await ops_test.model.relate("mlflow-server:dashboard-links", "kubeflow-dashboard:links")
@@ -72,15 +87,20 @@ async def test_bundle_deployment_works(self, ops_test: OpsTest, lightkube_client
         assert "Email Address" in result_text
         assert "Password" in result_text
 
+
 def get_public_url(lightkube_client: lightkube.Client, bundle_name: str):
     """Extracts public URL from service istio-ingressgateway-workload."""
     ingressgateway_svc = lightkube_client.get(
         Service, "istio-ingressgateway-workload", namespace=bundle_name
     )
-    address = ingressgateway_svc.status.loadBalancer.ingress[0].hostname or ingressgateway_svc.status.loadBalancer.ingress[0].ip
+    address = (
+        ingressgateway_svc.status.loadBalancer.ingress[0].hostname
+        or ingressgateway_svc.status.loadBalancer.ingress[0].ip
+    )
     public_url = f"http://{address}"
     return public_url
 
+
 async def fetch_response(url, headers=None):
     """Fetch provided URL and return (status, text)."""
     result_status = 0