filter online devices before fail fast

OpenTMI · Aug 1, 2024 · c5a5217 · c5a5217
1 parent a608afc
commit c5a5217
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/stf_appium_client/StfClient.py b/stf_appium_client/StfClient.py
@@ -14,6 +14,8 @@
 from stf_client.api.user_api import UserApi
 from stf_client.api.devices_api import DevicesApi
 
+STATUS_ONLINE = 3
+
 
 class StfClient(Logger):
     DEFAULT_ALLOCATION_TIMEOUT_SECONDS = 900
@@ -173,7 +175,7 @@ def list_devices(self, requirements: dict, fields: str = "", available_filter: b
                     ready=True,
                     using=False,
                     owner=None,
-                    status=3)  # 3=Online
+                    status=STATUS_ONLINE)
             )
 
         self.logger.debug(
@@ -183,6 +185,12 @@ def list_devices(self, requirements: dict, fields: str = "", available_filter: b
 
         return filter_(devices, predicate)
 
+    def list_online_devices(self, requirements: dict):
+        suitable_devices = self.list_devices(requirements=requirements)
+        online = filter(lambda device: device.get('status') == STATUS_ONLINE, suitable_devices)
+        return list(online)
+
+
     def find_and_allocate(self, requirements: dict,
                           timeout_seconds: int = DEFAULT_ALLOCATION_TIMEOUT_SECONDS,
                           shuffle: bool = True) -> dict:
@@ -198,7 +206,7 @@ def find_and_allocate(self, requirements: dict,
         :raises DeviceNotFound: suitable device not found or all devices are allocated already
         """
         NotConnectedError.invariant(self._client, 'Not connected')
-        suitable_devices = self.list_devices(requirements=requirements)
+        suitable_devices = self.list_online_devices(requirements=requirements)
         DeviceNotFound.invariant(len(suitable_devices), 'no suitable devices found')
         if shuffle:
             random.shuffle(suitable_devices)
@@ -239,7 +247,7 @@ def find_wait_and_allocate(self,
         suitable_devices = self.list_devices(requirements=requirements, available_filter=False)
         if not suitable_devices:
             raise DeviceNotFound(f'No suitable devices found ({json.dumps(requirements)})')
-        
+
         print(f'wait_until: {wait_until}')
         while True:
             remaining_time = int(wait_until - time.time())