trigger updated

astronomer · Jun 13, 2024 · fe0b405 · fe0b405
1 parent 8366cad
commit fe0b405
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/tests/triggers/test_anyscale_triggers.py b/tests/triggers/test_anyscale_triggers.py
@@ -72,6 +72,7 @@ def test_get_current_status(self, mock_get_job_status):
 
             # Call the method to test
             status = trigger.get_current_status('123')
+
 
             # Verify the result
             self.assertEqual(status, 'SUCCEEDED')
@@ -129,9 +130,8 @@ def test_serialize(self):
         # Check if the result matches the expected output
         self.assertEqual(result, expected_output)
 
-    @pytest.mark.asyncio
-    @mock.patch("anyscale_provider.hooks.anyscale.AnyscaleHook.get_job_status")
-    @mock.patch("anyscale_provider.hooks.anyscale.AnyscaleHook.get_logs")
+    @patch("anyscale_provider.hooks.anyscale.AnyscaleHook.get_job_status")
+    @patch("anyscale_provider.hooks.anyscale.AnyscaleHook.get_logs")
     async def test_anyscale_run_trigger(self, mocked_get_logs, mocked_get_job_status):
         """Test AnyscaleJobTrigger run method with mocked details."""
         mocked_get_job_status.return_value.state = JobState.SUCCEEDED
@@ -140,22 +140,22 @@ async def test_anyscale_run_trigger(self, mocked_get_logs, mocked_get_job_status
         trigger = AnyscaleJobTrigger(
             conn_id="test_conn",
             job_id="1234",
-            job_start_time=time.time(),
+            job_start_time=datetime.now().timestamp(),
             poll_interval=1,
             timeout=5,
         )
 
         task = asyncio.create_task(trigger.run().__anext__())
         await asyncio.sleep(0.5)
 
-        assert task.done() is False
+        self.assertFalse(task.done())
 
         await asyncio.sleep(2)
         result = await task
 
-        assert result.payload["status"] == JobState.SUCCEEDED
-        assert result.payload["message"] == "Job 1234 completed with status JobState.SUCCEEDED."
-        assert result.payload["job_id"] == "1234"
+        self.assertEqual(result["status"], JobState.SUCCEEDED)
+        self.assertEqual(result["message"], "Job 1234 completed with status JobState.SUCCEEDED.")
+        self.assertEqual(result["job_id"], "1234")
 
 
 class TestAnyscaleServiceTrigger(unittest.TestCase):