fixed imports in mocks (#12146)

GitOrigin-RevId: 3298705c0e49b6312f73ee072997acd624ad48f1
descarteslabs · Sep 8, 2023 · c3a9024 · c3a9024
1 parent cbad043
commit c3a9024
Showing 1 changed file with 14 additions and 25 deletions.
diff --git a/descarteslabs/core/compute/tests/test_job.py b/descarteslabs/core/compute/tests/test_job.py
@@ -1,9 +1,10 @@
 import json
 from datetime import timezone
-from unittest import mock
+from unittest.mock import patch
 
 import responses
 
+from descarteslabs.catalog import Blob
 from descarteslabs.compute import Job, JobStatus
 
 from .base import BaseTestCase
@@ -160,46 +161,34 @@ def test_delete_failed(self):
         assert job.state == "saved"
 
     @responses.activate
-    @mock.patch(
-        "descarteslabs.compute.job.Job._get_result_namespace",
-        return_value="some-org:some-user",
-    )
-    @mock.patch(
-        "descarteslabs.catalog.blob.Blob.get_data",
-        return_value=json.dumps({"test": "blah"}).encode(),
-    )
+    @patch.object(Job, "_get_result_namespace")
+    @patch.object(Blob, "get_data")
     def test_result_json(self, mock_get_data, mock_get_result_namespace):
+        mock_get_data.return_value = {"test": "blah"}
+        mock_get_result_namespace.return_value = ("some-org:some-user",)
         job = Job(id="some-id", function_id="some-fn", saved=True)
         assert job.result() == {"test": "blah"}
 
     @responses.activate
-    @mock.patch(
-        "descarteslabs.compute.job.Job._get_result_namespace",
-        return_value="some-org:some-user",
-    )
-    @mock.patch(
-        "descarteslabs.catalog.blob.Blob.get_data",
-        return_value=json.dumps(15.68).encode(),
-    )
+    @patch.object(Job, "_get_result_namespace")
+    @patch.object(Blob, "get_data")
     def test_result_float(self, mock_get_data, mock_get_result_namespace):
+        mock_get_data.return_value = json.dumps(15.68).encode()
+        mock_get_result_namespace.return_value = ("some-org:some-user",)
         job = Job(id="some-id", function_id="some-fn", saved=True)
         assert job.result() == 15.68
 
     @responses.activate
-    @mock.patch(
-        "descarteslabs.compute.job.Job._get_result_namespace",
-        return_value="some-org:some-user",
-    )
-    @mock.patch(
-        "descarteslabs.catalog.blob.Blob.get_data",
-        return_value="blah",
-    )
+    @patch.object(Job, "_get_result_namespace")
+    @patch.object(Blob, "get_data")
     def test_result_cast(self, mock_get_data, mock_get_result_namespace):
         class CustomString:
             @classmethod
             def deserialize(cls, data: bytes):
                 return "custom"
 
+        mock_get_data.return_value = "blah"
+        mock_get_result_namespace.return_value = ("some-org:some-user",)
         job = Job(id="some-id", function_id="some-fn", saved=True)
         assert job.result(CustomString) == "custom"