try once even wait_until is zero and improve unittests

OpenTMI · Mar 23, 2024 · 118690a · 118690a
1 parent 0243e07
commit 118690a
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/stf_appium_client/StfClient.py b/stf_appium_client/StfClient.py
@@ -233,8 +233,10 @@ def find_wait_and_allocate(self,
         :return: device dictionary
         """
         wait_until = time.time() + wait_timeout
-        while wait_until - time.time() > 0:
+        print(f'wait_until: {wait_until}')
+        while True:
             remaining_time = int(wait_until - time.time())
+            print(f'remaining_time: {remaining_time}')
             try:
                 return self.find_and_allocate(requirements=requirements,
                                               timeout_seconds=timeout_seconds,
@@ -243,6 +245,8 @@ def find_wait_and_allocate(self,
                 # Wait a while
                 self.logger.debug(f'Suitable device not available, '
                                   f'wait a while and try again. Timeout in {remaining_time} seconds')
+            if (wait_until - time.time()) <= 0:
+                break
             # Wait a while to avoid too frequent polling
             time.sleep(1)
         raise DeviceNotFound(f'Suitable device not found within {wait_timeout}s timeout ({json.dumps(requirements)})')

diff --git a/test/test_StfClient.py b/test/test_StfClient.py
@@ -222,3 +222,14 @@ def test_allocation_context_timeout(self, mock_sleep):
                 pass
         self.assertEqual(str(error.exception), 'Suitable device not found within 0s timeout ({"serial": "123"})')
 
+    @patch('time.sleep', return_value=MagicMock())
+    @patch('time.time')
+    def test_allocation_context_timeout_long(self, mock_time, mock_sleep):
+        dev1 = {'serial': '123', 'present': True, 'ready': True, 'using': True, 'owner': "asd", 'status': 3}
+        self.client.get_devices = MagicMock(return_value=[dev1])
+        self.client.stf_find_and_allocate = MagicMock(side_effect=DeviceNotFound)
+        mock_time.side_effect = [0, 0, 0, 10, 10, 10]
+        with self.assertRaises(DeviceNotFound) as error:
+            with self.client.allocation_context({"serial": '123'}, wait_timeout=10):
+                pass
+        self.assertEqual(str(error.exception), 'Suitable device not found within 10s timeout ({"serial": "123"})')