runtime(lava): Add methods to get queue size and all device names for device type

Signed-off-by: Denys Fedoryshchenko <denys.f@collabora.com>

Author: nuclearcat

kernelci/runtime/lava.py

@@ -403,6 +403,100 @@ def _connect(self):
         }
         return rest_api
 
+    def _get_response(self, url, params=None):
+        resp = self._server.session.get(url, params=params, timeout=30)
+        resp.raise_for_status()
+        return resp.json()
+
+    def _get_all(self, url, params=None):
+        resp = self._get_response(url, params=params)
+        results = resp.get('results', [])
+        next_url = resp.get('next')
+        while next_url:
+            resp = self._get_response(next_url)
+            results.extend(resp.get('results', []))
+            next_url = resp.get('next')
+        return results
+
+    def get_devicetype_job_count(self, device_types):
+        """Get queued job counts per requested device type.
+
+        *device_types* can be a device type name or list of device type names.
+        This uses scheduler.jobs.queue which is the only reliable way to query
+        queue sizes per device type.
+        """
+        if self._server.url is None:
+            raise ValueError("LAVA server URL is not configured")
+
+        single_type = isinstance(device_types, str)
+        if single_type:
+            requested_types = [device_types]
+        else:
+            requested_types = list(device_types or [])
+        if not requested_types:
+            return 0 if single_type else {}
+
+        queue_url = urljoin(self._server.url, 'scheduler/jobs/queue/')
+        counts = {device_type: 0 for device_type in requested_types}
+        start = 0
+        limit = 100
+        while True:
+            params = {
+                'device_types': requested_types,
+                'start': start,
+                'limit': limit,
+            }
+            jobs = self._get_response(queue_url, params=params)
+            for job in jobs:
+                requested_type = job.get('requested_device_type')
+                if requested_type in counts:
+                    counts[requested_type] += 1
+            if len(jobs) < limit:
+                break
+            start += limit
+
+        if single_type:
+            return counts.get(requested_types[0], 0)
+        return counts
+
+    def get_device_names_by_type(self, device_type, online_only=False):
+        """Get device names for a given LAVA device type.
+
+        *device_type* can be a string or list of device type names.
+        *online_only* filters devices with health == 'Good' when available.
+        Use this with get_devicetype_job_count() to gate submissions when the
+        queue per device type exceeds a threshold.
+        """
+        if self._server.url is None:
+            raise ValueError("LAVA server URL is not configured")
+
+        single_type = isinstance(device_type, str)
+        if single_type:
+            device_types = [device_type]
+        else:
+            device_types = list(device_type or [])
+        if not device_types:
+            return [] if single_type else {}
+
+        devices_url = urljoin(self._server.url, 'devices/')
+        result = {}
+        for dev_type in device_types:
+            params = {'device_type': dev_type}
+            devices = self._get_all(devices_url, params=params)
+            names = []
+            for device in devices:
+                if device.get('device_type') != dev_type:
+                    continue
+                if online_only and device.get('health') not in (None, 'Good'):
+                    continue
+                hostname = device.get('hostname') or device.get('name')
+                if hostname:
+                    names.append(hostname)
+            result[dev_type] = names
+        if single_type:
+            return result.get(device_types[0], [])
+        return result
+
     def _submit(self, job):
         if self._server.url is None:
             return self._store_job_in_external_storage(job)

tests/test_runtime.py

@@ -12,10 +12,53 @@
 # implementation.
 # pylint: disable=protected-access
 
+import types
+
 import kernelci.config
 import kernelci.runtime
 
 
+class _FakeResponse:
+    def __init__(self, payload, status_code=200, text=""):
+        """Initialize a fake response with payload and status."""
+        self._payload = payload
+        self.status_code = status_code
+        self.text = text
+
+    def raise_for_status(self):
+        """Raise an error when the response indicates failure."""
+        if self.status_code >= 400:
+            raise RuntimeError(f"HTTP {self.status_code}")
+
+    def json(self):
+        """Return the preloaded JSON payload."""
+        return self._payload
+
+
+class _FakeSession:
+    def __init__(self, get_handler=None, post_handler=None):
+        """Initialize a fake session with optional handlers."""
+        self._get_handler = get_handler
+        self._post_handler = post_handler
+        self.calls = []
+
+    def get(self, url, params=None, timeout=30):  # pylint: disable=unused-argument
+        """Invoke the GET handler and return a fake response."""
+        if not self._get_handler:
+            raise AssertionError("GET handler not set")
+        self.calls.append((url, params))
+        return _FakeResponse(self._get_handler(url, params))
+
+    def post(  # pylint: disable=unused-argument
+        self, url, json=None, allow_redirects=False, timeout=30
+    ):
+        """Invoke the POST handler and return its response."""
+        if not self._post_handler:
+            raise AssertionError("POST handler not set")
+        self.calls.append((url, json))
+        return self._post_handler(url, json)
+
+
 def test_runtimes_init():
     """Test that all the runtimes can be initialised (offline)"""
     config = kernelci.config.load('tests/configs/runtimes.yaml')
@@ -66,3 +109,94 @@ def test_lava_priority_scale():
             spec_priority = int(priority)
             print(f"* {plan_name:12s} {lab_priority:3d} {spec_priority:3d}")
             assert lab_priority == spec_priority
+
+
+def test_lava_get_devicetype_job_count():
+    """Test queued job count aggregation via scheduler.jobs.queue."""
+    config = kernelci.config.load('tests/configs/lava-runtimes.yaml')
+    runtime_config = config['runtimes']['lab-min-12-max-40-new-runtime']
+    lab = kernelci.runtime.get_runtime(runtime_config)
+
+    def handler(url, params):
+        assert url.endswith('scheduler/jobs/queue/')
+        assert params.get('device_types') == ['type-a', 'type-b']
+        assert params.get('limit') == 100
+        if params.get('start') == 0:
+            jobs = (
+                [{'requested_device_type': 'type-a'} for _ in range(60)] +
+                [{'requested_device_type': 'type-b'} for _ in range(40)]
+            )
+            return jobs
+        if params.get('start') == 100:
+            return [{'requested_device_type': 'type-a'}]
+        raise AssertionError(f"Unexpected request: {url} {params}")
+
+    lab._server = types.SimpleNamespace(
+        url='http://lava/api/v0.2/',
+        session=_FakeSession(get_handler=handler),
+    )
+
+    counts = lab.get_devicetype_job_count(['type-a', 'type-b'])
+    assert counts == {'type-a': 61, 'type-b': 40}
+
+
+def test_lava_get_device_names_by_type():
+    """Test device name lookups with type filtering and health checks."""
+    config = kernelci.config.load('tests/configs/lava-runtimes.yaml')
+    runtime_config = config['runtimes']['lab-min-12-max-40-new-runtime']
+    lab = kernelci.runtime.get_runtime(runtime_config)
+
+    def handler(url, params):
+        if url.endswith('devices/'):
+            dev_type = params.get('device_type')
+            if dev_type == 'type-a':
+                return {
+                    'results': [
+                        {'device_type': 'type-a', 'hostname': 'dev-1', 'health': 'Good'},
+                        {'device_type': 'type-a', 'hostname': 'dev-2', 'health': 'Bad'},
+                        {'device_type': 'type-b', 'hostname': 'dev-x', 'health': 'Good'},
+                    ],
+                    'next': None,
+                }
+            if dev_type == 'type-b':
+                return {
+                    'results': [
+                        {'device_type': 'type-b', 'hostname': 'dev-3', 'health': 'Good'},
+                    ],
+                    'next': None,
+                }
+        raise AssertionError(f"Unexpected request: {url} {params}")
+
+    lab._server = types.SimpleNamespace(
+        url='http://lava/api/v0.2/',
+        session=_FakeSession(get_handler=handler),
+    )
+
+    names = lab.get_device_names_by_type('type-a', online_only=True)
+    assert names == ['dev-1']
+
+    names_by_type = lab.get_device_names_by_type(['type-a', 'type-b'])
+    assert names_by_type == {'type-a': ['dev-1', 'dev-2'], 'type-b': ['dev-3']}
+
+
+def test_lava_submit_rest():
+    """Test LAVA REST submission builds a job with expected payload."""
+    config = kernelci.config.load('tests/configs/lava-runtimes.yaml')
+    runtime_config = config['runtimes']['lab-min-12-max-40-new-runtime']
+    lab = kernelci.runtime.get_runtime(runtime_config)
+
+    captured = {}
+
+    def post_handler(url, payload):
+        assert url.endswith('jobs/')
+        captured['json'] = payload
+        return _FakeResponse({'job_ids': [123]})
+
+    lab._server = types.SimpleNamespace(
+        url='http://lava/api/v0.2/',
+        session=_FakeSession(post_handler=post_handler),
+    )
+
+    job_id = lab._submit("jobdef")
+    assert job_id == 123
+    assert captured['json']['definition'] == "jobdef"