improve unittest and handle ForbiddenException properly

stf_appium_client/StfClient.py

@@ -207,7 +207,7 @@ def find_and_allocate(self, requirements: dict,
         def try_allocate(device_candidate):
             try:
                 return self.allocate(device_candidate, timeout_seconds=timeout_seconds)
-            except AssertionError as error:
+            except (AssertionError, ForbiddenException) as error:
                 self.logger.warning(f"{device_candidate.get('serial')} allocation fails: {error}")
                 return None
 
@@ -245,7 +245,7 @@ def find_wait_and_allocate(self,
                                   f'wait a while and try again. Timeout in {remaining_time} seconds')
             # Wait a while to avoid too frequent polling
             time.sleep(1)
-        raise DeviceNotFound(f'Suitable device not found within timeout ({wait_timeout})')
+        raise DeviceNotFound(f'Suitable device not found within {wait_timeout}s timeout ({json.dumps(requirements)})')
 
     @contextmanager
     def allocation_context(self, requirements: dict,

test/test_StfClient.py

@@ -9,8 +9,6 @@
 from stf_appium_client.exceptions import *
 
 
-
-
 class TestStfClientBasics(unittest.TestCase):
 
     @classmethod
@@ -214,4 +212,13 @@ def test_allocation_context_wait_success(self, mock_sleep):
             self.assertEqual(device['serial'], '123')
             self.assertEqual(device['remote_adb_url'], url)
 
+    @patch('time.sleep', side_effect=MagicMock())
+    def test_allocation_context_timeout(self, mock_sleep):
+        dev1 = {'serial': '123', 'present': True, 'ready': True, 'using': True, 'owner': "asd", 'status': 3}
+        self.client.get_devices = MagicMock(return_value=[dev1])
+
+        with self.assertRaises(DeviceNotFound) as error:
+            with self.client.allocation_context({"serial": '123'}, wait_timeout=0) as device:
+                pass
+        self.assertEqual(str(error.exception), 'Suitable device not found within 0s timeout ({"serial": "123"})')