Add tests

Author: VipulMascarenhas

tests/unitary/with_extras/aqua/test_model.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*--
-# Copyright (c) 2024 Oracle and/or its affiliates.
+# Copyright (c) 2024, 2025 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
 import os
@@ -653,12 +653,12 @@ def test_get_model_fine_tuned(
         }
 
     @pytest.mark.parametrize(
-        ("artifact_location_set", "download_from_hf"),
+        ("artifact_location_set", "download_from_hf", "delete_from_local"),
         [
-            (True, True),
-            (True, False),
-            (False, True),
-            (False, False),
+            (True, True, True),
+            (True, False, True),
+            (False, True, False),
+            (False, False, True),
         ],
     )
     @patch("ads.model.service.oci_datascience_model.OCIDataScienceModel.create")
@@ -683,6 +683,7 @@ def test_import_verified_model(
         mock_ocidsc_create,
         artifact_location_set,
         download_from_hf,
+        delete_from_local,
         mock_get_hf_model_info,
         mock_init_client,
     ):
@@ -747,6 +748,7 @@ def test_import_verified_model(
                     os_path=os_path,
                     local_dir=str(tmpdir),
                     download_from_hf=True,
+                    delete_from_local=delete_from_local,
                     allow_patterns=["*.json"],
                     ignore_patterns=["test.json"],
                 )
@@ -761,6 +763,20 @@ def test_import_verified_model(
                         f"oci os object bulk-upload --src-dir {str(tmpdir)}/{model_name} --prefix prefix/path/{model_name}/ -bn aqua-bkt -ns aqua-ns --auth api_key --profile DEFAULT --no-overwrite --exclude {HF_METADATA_FOLDER}*"
                     )
                 )
+                if delete_from_local:
+                    cache_dir = os.path.join(
+                        os.path.expanduser("~"),
+                        ".cache",
+                        "huggingface",
+                        "hub",
+                        "models--oracle--aqua-1t-mega-model",
+                    )
+                    assert (
+                        os.path.exists(f"{str(tmpdir)}/{os.path.dirname(model_name)}")
+                        is False
+                    )
+                    assert os.path.exists(cache_dir) is False
+
         else:
             model: AquaModel = app.register(
                 model="ocid1.datasciencemodel.xxx.xxxx.",
@@ -1183,22 +1199,22 @@ def test_import_model_with_input_tags(
                     "model": "oracle/oracle-1it",
                     "inference_container": "odsc-vllm-serving",
                 },
-                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --inference_container odsc-vllm-serving",
+                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --delete_from_local True --inference_container odsc-vllm-serving",
             ),
             (
                 {
                     "os_path": "oci://aqua-bkt@aqua-ns/path",
                     "model": "ocid1.datasciencemodel.oc1.iad.<OCID>",
                 },
-                "ads aqua model register --model ocid1.datasciencemodel.oc1.iad.<OCID> --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True",
+                "ads aqua model register --model ocid1.datasciencemodel.oc1.iad.<OCID> --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --delete_from_local True",
             ),
             (
                 {
                     "os_path": "oci://aqua-bkt@aqua-ns/path",
                     "model": "oracle/oracle-1it",
                     "download_from_hf": False,
                 },
-                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf False",
+                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf False --delete_from_local True",
             ),
             (
                 {
@@ -1207,7 +1223,7 @@ def test_import_model_with_input_tags(
                     "download_from_hf": True,
                     "model_file": "test_model_file",
                 },
-                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --model_file test_model_file",
+                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --delete_from_local True --model_file test_model_file",
             ),
             (
                 {
@@ -1216,7 +1232,7 @@ def test_import_model_with_input_tags(
                     "inference_container": "odsc-tei-serving",
                     "inference_container_uri": "<region>.ocir.io/<your_tenancy>/<your_image>",
                 },
-                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --inference_container odsc-tei-serving --inference_container_uri <region>.ocir.io/<your_tenancy>/<your_image>",
+                "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path --download_from_hf True --delete_from_local True --inference_container odsc-tei-serving --inference_container_uri <region>.ocir.io/<your_tenancy>/<your_image>",
             ),
             (
                 {
@@ -1227,7 +1243,7 @@ def test_import_model_with_input_tags(
                     "defined_tags": {"dtag1": "dvalue1", "dtag2": "dvalue2"},
                 },
                 "ads aqua model register --model oracle/oracle-1it --os_path oci://aqua-bkt@aqua-ns/path "
-                "--download_from_hf True --inference_container odsc-vllm-serving --freeform_tags "
+                "--download_from_hf True --delete_from_local True --inference_container odsc-vllm-serving --freeform_tags "
                 '{"ftag1": "fvalue1", "ftag2": "fvalue2"} --defined_tags {"dtag1": "dvalue1", "dtag2": "dvalue2"}',
             ),
         ],