amend time-machine update in aws provider test and task-sdk tests

Author: bugraoz93

providers/amazon/tests/unit/amazon/aws/sensors/test_s3.py

@@ -556,9 +556,9 @@ def test_poke_succeeds_on_upload_complete(self, time_machine):
         self.sensor.hook = mock.MagicMock()
         self.sensor.hook.list_keys.return_value = {"a"}
         assert not self.sensor.poke(dict())
-        time_machine.coordinates.shift(10)
+        time_machine.shift(10)
         assert not self.sensor.poke(dict())
-        time_machine.coordinates.shift(10)
+        time_machine.shift(10)
         assert self.sensor.poke(dict())
 
     def test_fail_is_keys_unchanged(self):

task-sdk/tests/task_sdk/bases/test_sensor.py

@@ -175,15 +175,15 @@ def test_ok_with_reschedule(self, run_task, make_sensor, time_machine):
         assert msg.reschedule_date == date1 + timedelta(seconds=sensor.poke_interval)
 
         # second poke returns False and task is re-scheduled
-        time_machine.coordinates.shift(sensor.poke_interval)
+        time_machine.shift(sensor.poke_interval)
         date2 = date1 + timedelta(seconds=sensor.poke_interval)
         state, msg, _ = run_task(task=sensor)
 
         assert state == TaskInstanceState.UP_FOR_RESCHEDULE
         assert msg.reschedule_date == date2 + timedelta(seconds=sensor.poke_interval)
 
         # third poke returns True and task succeeds
-        time_machine.coordinates.shift(sensor.poke_interval)
+        time_machine.shift(sensor.poke_interval)
         state, _, _ = run_task(task=sensor)
 
         assert state == TaskInstanceState.SUCCESS
@@ -201,7 +201,7 @@ def test_fail_with_reschedule(self, run_task, make_sensor, time_machine, mock_su
         assert msg.reschedule_date == date1 + timedelta(seconds=sensor.poke_interval)
 
         # second poke returns False, timeout occurs
-        time_machine.coordinates.shift(sensor.poke_interval)
+        time_machine.shift(sensor.poke_interval)
 
         # Mocking values from DB/API-server
         mock_supervisor_comms.send.return_value = TaskRescheduleStartDate(start_date=date1)
@@ -223,7 +223,7 @@ def test_soft_fail_with_reschedule(self, run_task, make_sensor, time_machine, mo
         assert state == TaskInstanceState.UP_FOR_RESCHEDULE
 
         # second poke returns False, timeout occurs
-        time_machine.coordinates.shift(sensor.poke_interval)
+        time_machine.shift(sensor.poke_interval)
 
         # Mocking values from DB/API-server
         mock_supervisor_comms.send.return_value = TaskRescheduleStartDate(start_date=date1)
@@ -258,7 +258,7 @@ def run_duration():
         # loop poke returns false
         for _poke_count in range(1, false_count + 1):
             curr_date = curr_date + timedelta(seconds=new_interval)
-            time_machine.coordinates.shift(new_interval)
+            time_machine.shift(new_interval)
             state, msg, _ = run_task(sensor, context_update={"task_reschedule_count": _poke_count})
             assert state == TaskInstanceState.UP_FOR_RESCHEDULE
             old_interval = new_interval
@@ -268,7 +268,7 @@ def run_duration():
 
         # last poke returns True and task succeeds
         curr_date = curr_date + timedelta(seconds=new_interval)
-        time_machine.coordinates.shift(new_interval)
+        time_machine.shift(new_interval)
 
         state, msg, _ = run_task(sensor, context_update={"task_reschedule_count": false_count + 1})
         assert state == TaskInstanceState.SUCCESS