Added unit tests for model and ft handlers (#774)

Author: lu-ohai

tests/unitary/with_extras/aqua/test_finetuning_handler.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*--
+
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+from unittest import TestCase
+from unittest.mock import MagicMock
+from mock import patch
+
+from notebook.base.handlers import IPythonHandler
+from ads.aqua.extension.finetune_handler import AquaFineTuneHandler
+from ads.aqua.finetune import AquaFineTuningApp, CreateFineTuningDetails
+
+
+class TestDataset:
+    mock_valid_input = dict(
+        ft_source_id="ocid1.datasciencemodel.oc1.iad.<OCID>",
+        ft_name="test_ft_name",
+        dataset_path="oci://ds_bucket@namespace/prefix/dataset.jsonl",
+        report_path="oci://report_bucket@namespace/prefix/",
+        ft_parameters={
+            "epochs":1,
+            "learning_rate":0.02
+        },
+        shape_name="VM.GPU.A10.1",
+        replica=1,
+        validation_set_size=0.2,
+        block_storage_size=1,
+        experiment_name="test_experiment_name",
+    )
+
+    mock_finetuning_config = {
+        "shape": {
+            "VM.GPU.A10.1": {
+                "batch_size": 1,
+                "replica": 1
+            },
+        }
+    }
+
+
+class FineTuningHandlerTestCase(TestCase):
+
+    @patch.object(IPythonHandler, "__init__")
+    def setUp(self, ipython_init_mock) -> None:
+        ipython_init_mock.return_value = None
+        self.test_instance = AquaFineTuneHandler(MagicMock(), MagicMock())
+        self.test_instance.request = MagicMock()
+        self.test_instance.finish = MagicMock()
+
+    @patch.object(AquaFineTuneHandler, "get_finetuning_config")
+    @patch("ads.aqua.extension.finetune_handler.urlparse")
+    def test_get(self, mock_urlparse, mock_get_finetuning_config):
+        request_path = MagicMock(path="aqua/finetuning/config")
+        mock_urlparse.return_value = request_path
+
+        mock_get_finetuning_config.return_value = TestDataset.mock_finetuning_config
+
+        fineruning_config = self.test_instance.get(id="test_model_id")
+        mock_urlparse.assert_called()
+        mock_get_finetuning_config.assert_called_with("test_model_id")
+        assert fineruning_config == TestDataset.mock_finetuning_config
+
+    @patch.object(AquaFineTuningApp, "create")
+    def test_post(self, mock_create):
+        self.test_instance.get_json_body = MagicMock(
+            return_value = TestDataset.mock_valid_input
+        )
+        self.test_instance.post()
+
+        self.test_instance.finish.assert_called_with(
+            mock_create.return_value
+        )
+        mock_create.assert_called_with(
+            CreateFineTuningDetails(**TestDataset.mock_valid_input)
+        )
+
+    @patch.object(AquaFineTuningApp, "get_finetuning_config")
+    def test_get_finetuning_config(self, mock_get_finetuning_config):
+        mock_get_finetuning_config.return_value = TestDataset.mock_finetuning_config
+
+        self.test_instance.get_finetuning_config(model_id="test_model_id")
+
+        self.test_instance.finish.assert_called_with(
+            TestDataset.mock_finetuning_config
+        )
+        mock_get_finetuning_config.assert_called_with(
+            model_id="test_model_id"
+        )

tests/unitary/with_extras/aqua/test_model_handler.py

@@ -0,0 +1,82 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*--
+
+# Copyright (c) 2024 Oracle and/or its affiliates.
+# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
+
+from unittest import TestCase
+from unittest.mock import MagicMock
+
+from mock import patch
+from notebook.base.handlers import IPythonHandler
+
+from ads.aqua.extension.model_handler import AquaModelHandler, AquaModelLicenseHandler
+from ads.aqua.model import AquaModelApp
+
+
+class ModelHandlerTestCase(TestCase):
+    
+    @patch.object(IPythonHandler, "__init__")
+    def setUp(self, ipython_init_mock) -> None:
+        ipython_init_mock.return_value = None
+        self.model_handler = AquaModelHandler(MagicMock(), MagicMock())
+        self.model_handler.request = MagicMock()
+        self.model_handler.finish = MagicMock()
+
+    @patch.object(AquaModelHandler, "list")
+    def test_get_no_id(self, mock_list):
+        self.model_handler.get()
+        mock_list.assert_called()
+
+    @patch.object(AquaModelHandler, "read")
+    def test_get_with_id(self, mock_read):
+        self.model_handler.get(model_id="test_model_id")
+        mock_read.assert_called_with("test_model_id")
+
+    @patch.object(AquaModelApp, "get")
+    def test_read(self, mock_get):
+        self.model_handler.read(model_id="test_model_id")
+        self.model_handler.finish.assert_called_with(
+            mock_get.return_value
+        )
+        mock_get.assert_called_with("test_model_id")
+
+    @patch.object(AquaModelApp, "clear_model_list_cache")
+    @patch("ads.aqua.extension.model_handler.urlparse")
+    def test_delete(self, mock_urlparse, mock_clear_model_list_cache):
+        request_path = MagicMock(path="aqua/model/cache")
+        mock_urlparse.return_value = request_path
+
+        self.model_handler.delete()
+        self.model_handler.finish.assert_called_with(
+            mock_clear_model_list_cache.return_value
+        )
+
+        mock_urlparse.assert_called()
+        mock_clear_model_list_cache.assert_called()
+
+    @patch.object(AquaModelApp, "list")
+    def test_list(self, mock_list):
+        self.model_handler.list()
+
+        self.model_handler.finish.assert_called_with(
+            mock_list.return_value
+        )
+        mock_list.assert_called_with(None, None)
+
+class ModelLicenseHandlerTestCase(TestCase):
+    
+    @patch.object(IPythonHandler, "__init__")
+    def setUp(self, ipython_init_mock) -> None:
+        ipython_init_mock.return_value = None
+        self.model_license_handler = AquaModelLicenseHandler(MagicMock(), MagicMock())
+        self.model_license_handler.finish = MagicMock()
+
+    @patch.object(AquaModelApp, "load_license")
+    def test_get(self, mock_load_license):
+        self.model_license_handler.get(model_id="test_model_id")
+
+        self.model_license_handler.finish.assert_called_with(
+            mock_load_license.return_value
+        )
+        mock_load_license.assert_called_with("test_model_id")