feat: add test for TrustyAI database migration (#372)

* feat: add test for TrustyAI database migration

* create PVC only TAI Svc.
* upload data to PVC
* patch CR to enable migration flags and DB + PVC config
* wait for migration
* patch to use DB only
* schedule metric collection request and verify metrics

* fix: remove unused fixture

* feat: remove unnecessary if statement

* fix: correct namespace

* Update tests/model_explainability/trustyai_service/service/test_trustyai_service.py

Co-authored-by: Adolfo Aguirrezabal <[email protected]>

* fix: remove unnecessary variables and patch

---------

Co-authored-by: Adolfo Aguirrezabal <[email protected]>

Author: sheltoncyril

tests/model_explainability/trustyai_service/conftest.py

@@ -291,7 +291,6 @@ def gaussian_credit_model(
     minio_service: Service,
     minio_data_connection: Secret,
     mlserver_runtime: ServingRuntime,
-    trustyai_service_with_pvc_storage: TrustyAIService,
     teardown_resources: bool,
 ) -> Generator[InferenceService, Any, Any]:
     gaussian_credit_model_kwargs = {

tests/model_explainability/trustyai_service/constants.py

@@ -6,7 +6,11 @@
 TAI_DATA_CONFIG: Dict[str, str] = {"filename": "data.csv", "format": "CSV"}
 TAI_METRICS_CONFIG: Dict[str, str] = {"schedule": "5s"}
 TAI_PVC_STORAGE_CONFIG: Dict[str, str] = {"format": "PVC", "folder": "/inputs", "size": "1Gi"}
-TAI_DB_STORAGE_CONFIG = {"format": "DATABASE", "size": "1Gi", "databaseConfigurations": "db-credentials"}
+TAI_DB_STORAGE_CONFIG: Dict[str, str] = {
+    "format": "DATABASE",
+    "size": "1Gi",
+    "databaseConfigurations": "db-credentials",
+}
 
 SKLEARN: str = "sklearn"
 MLSERVER: str = "mlserver"
@@ -56,3 +60,16 @@
 }
 
 ISVC_GETTER: str = "isvc-getter"
+
+TRUSTYAI_DB_MIGRATION_PATCH: dict[str, Any] = {
+    "metadata": {"annotations": {"trustyai.opendatahub.io/db-migration": "true"}},
+    "spec": {
+        "storage": {
+            "format": "DATABASE",
+            "folder": "/inputs",
+            "size": "1Gi",
+            "databaseConfigurations": "db-credentials",
+        },
+        "data": {"filename": "data.csv", "format": "BEAN"},
+    },
+}

tests/model_explainability/trustyai_service/service/test_trustyai_service.py

@@ -2,7 +2,10 @@
 from ocp_resources.namespace import Namespace
 from ocp_resources.trustyai_service import TrustyAIService
 
-from tests.model_explainability.trustyai_service.constants import DRIFT_BASE_DATA_PATH
+from tests.model_explainability.trustyai_service.constants import (
+    DRIFT_BASE_DATA_PATH,
+    TRUSTYAI_DB_MIGRATION_PATCH,
+)
 from tests.model_explainability.trustyai_service.trustyai_service_utils import (
     send_inferences_and_verify_trustyai_service_registered,
     verify_upload_data_to_trustyai_service,
@@ -14,6 +17,10 @@
     validate_trustyai_service_db_conn_failure,
     validate_trustyai_service_images,
 )
+from tests.model_explainability.trustyai_service.service.utils import (
+    wait_for_trustyai_db_migration_complete_log,
+    patch_trustyai_service_cr,
+)
 from utilities.constants import MinIo
 from utilities.manifests.openvino import OPENVINO_KSERVE_INFERENCE_CONFIG
 
@@ -176,3 +183,71 @@ def test_validate_trustyai_service_image(
         label_selector=f"app.kubernetes.io/instance={trustyai_service_with_pvc_storage.name}",
         trustyai_operator_configmap=trustyai_operator_configmap,
     )
+
+
+@pytest.mark.parametrize(
+    "model_namespace, minio_pod, minio_data_connection",
+    [
+        pytest.param(
+            {"name": "validate-trustyai-db-migration"},
+            MinIo.PodConfig.MODEL_MESH_MINIO_CONFIG,
+            {"bucket": MinIo.Buckets.MODELMESH_EXAMPLE_MODELS},
+        )
+    ],
+    indirect=True,
+)
+@pytest.mark.usefixtures("minio_pod")
+@pytest.mark.smoke
+def test_trustyai_service_db_migration(
+    admin_client,
+    current_client_token,
+    mariadb,
+    trustyai_db_ca_secret,
+    trustyai_service_with_pvc_storage,
+    gaussian_credit_model,
+) -> None:
+    """Verify if TrustyAI DB Migration works as expected.
+    This test initializes TrustyAI Service with PVC Storage at first with a database on standby but the service is not
+    configured to use it.
+    Data is uploaded to the PVC, then the TrustyAI CR is patched to trigger a migration from PVC to DB storage.
+    config.
+    Then waits for the migration success entry in the container logs and patches the service again to remove PVC config.
+    Finally, a metric is scheduled and checked if the service works as expected post migration.
+
+    Args:
+        admin_client: DynamicClient
+        current_client_token: RedactedString
+        mariadb: MariaDB
+        trustyai_db_ca_secret: None
+        trustyai_service_with_pvc_storage: TrustyAIService
+        gaussian_credit_model: Generator[InferenceService, Any, Any]
+
+    Returns:
+        None
+    """
+    verify_upload_data_to_trustyai_service(
+        client=admin_client,
+        trustyai_service=trustyai_service_with_pvc_storage,
+        token=current_client_token,
+        data_path=f"{DRIFT_BASE_DATA_PATH}/training_data.json",
+    )
+
+    trustyai_db_migration_patched_service = patch_trustyai_service_cr(
+        trustyai_service=trustyai_service_with_pvc_storage, patches=TRUSTYAI_DB_MIGRATION_PATCH
+    )
+
+    wait_for_trustyai_db_migration_complete_log(
+        client=admin_client,
+        trustyai_service=trustyai_db_migration_patched_service,
+    )
+
+    verify_trustyai_service_metric_scheduling_request(
+        client=admin_client,
+        trustyai_service=trustyai_db_migration_patched_service,
+        token=current_client_token,
+        metric_name=TrustyAIServiceMetrics.Drift.MEANSHIFT,
+        json_data={
+            "modelId": gaussian_credit_model.name,
+            "referenceTag": "TRAINING",
+        },
+    )

tests/model_explainability/trustyai_service/service/utils.py

@@ -0,0 +1,33 @@
+import re
+from typing import Any
+
+from kubernetes.dynamic import DynamicClient
+from ocp_resources.pod import Pod
+from ocp_resources.resource import ResourceEditor
+from ocp_resources.trustyai_service import TrustyAIService
+from timeout_sampler import retry
+
+from tests.model_explainability.trustyai_service.trustyai_service_utils import TRUSTYAI_SERVICE_NAME
+from utilities.constants import Timeout
+
+
+@retry(wait_timeout=Timeout.TIMEOUT_5MIN, sleep=5)
+def wait_for_trustyai_db_migration_complete_log(client: DynamicClient, trustyai_service: TrustyAIService) -> bool:
+    trustyai_pod = list(
+        Pod.get(
+            dyn_client=client,
+            namespace=trustyai_service.namespace,
+            label_selector=f"app.kubernetes.io/instance={trustyai_service.name}",
+        )
+    )[0]
+    return bool(
+        re.search(
+            r".+INFO.+Migration complete, the PVC is now safe to remove\.",
+            trustyai_pod.log(container=TRUSTYAI_SERVICE_NAME),
+        )
+    )
+
+
+def patch_trustyai_service_cr(trustyai_service: TrustyAIService, patches: dict[str, Any]) -> TrustyAIService:
+    ResourceEditor(patches={trustyai_service: patches}).update()
+    return trustyai_service