lint

Author: bpblanken

v03_pipeline/lib/tasks/trigger_hail_backend_reload_test.py

@@ -1,4 +1,4 @@
-from unittest.mock import Mock, patch
+from unittest import mock
 
 import hail as hl
 import luigi.worker
@@ -8,6 +8,7 @@
 from v03_pipeline.lib.model import DatasetType, ReferenceGenome, SampleType
 from v03_pipeline.lib.paths import variant_annotations_table_path
 from v03_pipeline.lib.tasks.trigger_hail_backend_reload import TriggerHailBackendReload
+from v03_pipeline.lib.test.mock_complete_task import MockCompleteTask
 from v03_pipeline.lib.test.mocked_dataroot_testcase import MockedDatarootTestCase
 
 TEST_VCF = 'v03_pipeline/var/test/callsets/1kg_30variants.vcf'
@@ -48,8 +49,16 @@ def setUp(self) -> None:
             ),
         )
 
-    @patch.object(requests, 'post')
-    def test_success(self, mock_post: Mock):
+    @mock.patch.object(requests, 'post')
+    @mock.patch(
+        'v03_pipeline.lib.tasks.trigger_hail_backend_reload.WriteSuccessFileTask',
+    )
+    def test_success(
+        self,
+        mock_write_success_file_task: mock.Mock,
+        mock_post: mock.Mock,
+    ):
+        mock_write_success_file_task.return_value = MockCompleteTask()
         mock_resp = requests.models.Response()
         mock_resp.status_code = 200
         mock_post.return_value = mock_resp
@@ -69,8 +78,14 @@ def test_success(self, mock_post: Mock):
         worker.run()
         self.assertTrue(task.complete())
 
-    @patch.object(requests, 'post')
-    def test_failure(self, mock_post: Mock):
+    @mock.patch.object(requests, 'post')
+    @mock.patch(
+        'v03_pipeline.lib.tasks.trigger_hail_backend_reload.WriteSuccessFileTask',
+    )
+    def test_failure(
+        self, mock_write_success_file_task: mock.Mock, mock_post: mock.Mock
+    ):
+        mock_write_success_file_task.return_value = MockCompleteTask()
         mock_resp = requests.models.Response()
         mock_resp.status_code = 500
         mock_post.return_value = mock_resp