|
20 | 20 | ) |
21 | 21 | from ads.aqua.model import AquaModelApp |
22 | 22 | from ads.aqua.model.entities import AquaModel, AquaModelSummary, HFModelSummary |
| 23 | +from ads.aqua.ui import AquaContainerConfig |
23 | 24 |
|
24 | 25 |
|
25 | 26 | class ModelHandlerTestCase(TestCase): |
@@ -93,6 +94,32 @@ def test_delete_with_id(self, mock_delete, mock_urlparse): |
93 | 94 | mock_urlparse.assert_called() |
94 | 95 | mock_delete.assert_called() |
95 | 96 |
|
| 97 | + @patch.object(AquaContainerConfig,"from_container_index_json") |
| 98 | + @patch.object(AquaModelApp,"edit_registered_model") |
| 99 | + def test_put(self,mock_edit,mock_container_index): |
| 100 | + mock_edit.return_value={"state":"EDITED"} |
| 101 | + mock_inference = MagicMock() |
| 102 | + mock_inference.values.return_value = [ |
| 103 | + MagicMock(family="odsc-vllm-serving"), |
| 104 | + MagicMock(family="odsc-tgi-serving"), |
| 105 | + MagicMock(family="odsc-vllm-serving"), |
| 106 | + ] |
| 107 | + |
| 108 | + mock_container_index.return_value = MagicMock(inference=mock_inference) |
| 109 | + self.model_handler.get_json_body = MagicMock( |
| 110 | + return_value=dict( |
| 111 | + task="text_generation", |
| 112 | + enable_finetuning="true", |
| 113 | + inference_container="odsc-tgi-serving", |
| 114 | + ) |
| 115 | + ) |
| 116 | + with patch("ads.aqua.extension.base_handler.AquaAPIhandler.finish") as mock_finish: |
| 117 | + mock_finish.side_effect = lambda x: x |
| 118 | + result = self.model_handler.put(id="ocid1.datasciencemodel.oc1.iad.xxx") |
| 119 | + print(f"result: ",result) |
| 120 | + assert result["state"] is "EDITED" |
| 121 | + mock_edit.assert_called() |
| 122 | + |
96 | 123 | @patch.object(AquaModelApp, "list") |
97 | 124 | def test_list(self, mock_list): |
98 | 125 | with patch( |
|
0 commit comments