Fixed unit test.

Author: lu-ohai

tests/unitary/with_extras/opctl/test_opctl_dataflow_backend.py

@@ -11,6 +11,7 @@
 
 import pytest
 import yaml
+from ads.jobs.builders.infrastructure.dataflow import DataFlowRun
 
 from ads.opctl.backend.ads_dataflow import DataFlowBackend
 
@@ -87,6 +88,23 @@ def test_dataflow_run(self, file_upload, job_run, job_create):
                 False,
             )
 
+    @patch(
+        "ads.opctl.backend.ads_dataflow.DataFlowRun.watch",
+        return_value=DataFlowRun(),
+    )
+    @patch(
+        "ads.opctl.backend.ads_dataflow.DataFlowRun.from_ocid",
+        return_value=DataFlowRun(),
+    )
+    def test_watch(self, mock_from_ocid, mock_watch):
+        config = self.config
+        config["execution"]["run_id"] = "test_dataflow_run_id"
+        config["execution"]["interval"] = 10
+        backend = DataFlowBackend(config)
+        backend.watch()
+        mock_from_ocid.assert_called_with("test_dataflow_run_id")
+        mock_watch.assert_called_with(interval=10)
+
     @pytest.mark.parametrize(
         "runtime_type",
         ["dataFlow", "dataFlowNotebook"],

tests/unitary/with_extras/opctl/test_opctl_ml_job_backend.py

@@ -109,6 +109,24 @@ def test_run_with_image(self, rt, job_run, job_create):
         job_create.assert_called()
         job_run.assert_called()
 
+    @patch(
+        "ads.opctl.backend.ads_ml_job.DataScienceJobRun.watch",
+        return_value=DataScienceJobRun(),
+    )
+    @patch(
+        "ads.opctl.backend.ads_ml_job.DataScienceJobRun.from_ocid",
+        return_value=DataScienceJobRun(),
+    )
+    def test_watch(self, mock_from_ocid, mock_watch):
+        config = self.config
+        config["execution"]["run_id"] = "test_job_run_id"
+        config["execution"]["interval"] = 10
+        config["execution"]["wait"] = 15
+        backend = MLJobBackend(config)
+        backend.watch()
+        mock_from_ocid.assert_called_with("test_job_run_id")
+        mock_watch.assert_called_with(interval=10, wait=15)
+
     @pytest.mark.parametrize(
         "runtime_type",
         ["container", "script", "python", "notebook", "gitPython"],

tests/unitary/with_extras/opctl/test_opctl_ml_pipeline_backend.py

@@ -129,10 +129,11 @@ def test_watch(self, mock_from_ocid, mock_watch):
         config = self.config
         config["execution"]["run_id"] = "test_pipeline_run_id"
         config["execution"]["log_type"] = "custom_log"
+        config["execution"]["interval"] = 10
         backend = PipelineBackend(config)
         backend.watch()
         mock_from_ocid.assert_called_with("test_pipeline_run_id")
-        mock_watch.assert_called_with(log_type="custom_log")
+        mock_watch.assert_called_with(interval=10, log_type="custom_log")
 
     @pytest.mark.parametrize(
         "runtime_type",