Use new set_instance_state

Author: NickCao

packages/jumpstarter-driver-corellium/jumpstarter_driver_corellium/corellium/api.py

@@ -103,26 +103,20 @@ def create_instance(self, name: str, project: Project, device: Device, os_versio
 
         return Instance(**data)
 
-    def get_instance(self, instance_ref: str) -> Optional[Instance]:
+    async def get_instance(self, instance_ref: str) -> Optional[Instance]:
         """
         Retrieve an existing instance by its name.
 
         Return None if it does not exist.
         """
-        try:
-            res = self.req.get(f"{self.baseurl}/v1/instances")
-            data = res.json()
-            res.raise_for_status()
-        except requests.exceptions.RequestException as e:
-            raise CorelliumApiException(data.get("error", str(e))) from e
-
-        for instance in data:
-            if instance["name"] == instance_ref or instance["id"] == instance_ref:
-                return Instance(id=instance["id"], state=instance["state"])
+        instances = await self.api.v1_get_instances()
+        for instance in instances:
+            if instance.name == instance_ref or instance.id == instance_ref:
+                return instance
 
         return None
 
-    def set_instance_state(self, instance: Instance, instance_state: str) -> None:
+    async def set_instance_state(self, instance: Instance, instance_state: str) -> None:
         """
         Set the virtual instance state from corellium.
 
@@ -138,16 +132,11 @@ def set_instance_state(self, instance: Instance, instance_state: str) -> None:
         - rebooting
         - error
         """
-        data = {"state": instance_state}
-
-        try:
-            res = self.req.put(f"{self.baseurl}/v1/instances/{instance.id}/state", json=data)
-            data = res.json() if res.status_code != 204 else None
-            res.raise_for_status()
-        except requests.exceptions.RequestException as e:
-            msgerr = data if data is not None else str(e)
 
-            raise CorelliumApiException(msgerr) from e
+        await self.api.v1_set_instance_state(
+            instance.id,
+            corellium_api.V1SetStateBody(state=instance_state),
+        )
 
     async def destroy_instance(self, instance: Instance) -> None:
         """

packages/jumpstarter-driver-corellium/jumpstarter_driver_corellium/driver.py

@@ -3,12 +3,12 @@
 """
 
 import os
-import time
 from collections.abc import AsyncGenerator
 from dataclasses import dataclass, field
 from datetime import datetime, timedelta
 from typing import Dict, Optional
 
+from anyio import sleep
 from jumpstarter_driver_power.driver import PowerReading, VirtualPowerInterface
 
 from .corellium.api import ApiClient
@@ -122,7 +122,7 @@ def get_timeout_opts(self) -> Dict[str, int]:
             "interval": os.environ.get("CORELLIUM_API_INTERVAL", 5),
         }
 
-    def wait_instance(self, current: Instance, desired: Optional[Instance]):
+    async def wait_instance(self, current: Instance, desired: Optional[Instance]):
         """
         Wait for `current` instance to reach the same state as the `desired` instance.
 
@@ -135,11 +135,11 @@ def wait_instance(self, current: Instance, desired: Optional[Instance]):
             if counter >= opts["retries"]:
                 raise ValueError(f"Instance took too long to be reach the desired state: {current}")
 
-            if self.parent.api.get_instance(current.id) == desired:
+            if await self.parent.api.get_instance(current.id) == desired:
                 break
 
             counter += 1
-            time.sleep(opts["interval"])
+            await sleep(opts["interval"])
 
     @export
     async def on(self) -> None:
@@ -164,9 +164,9 @@ async def on(self) -> None:
         self.logger.info(f"Using device spec: {device.name}")
 
         # retrieve an existing instance first
-        instance = self.parent.api.get_instance(self.parent.device_name)
+        instance = await self.parent.api.get_instance(self.parent.device_name)
         if instance:
-            self.parent.api.set_instance_state(instance, "on")
+            await self.parent.api.set_instance_state(instance, "on")
         # create a new one otherwise
         else:
             opts = {}
@@ -177,7 +177,7 @@ async def on(self) -> None:
             instance = self.parent.api.create_instance(self.parent.device_name, project, device, **opts)
         self.logger.info(f"Instance: {self.parent.device_name} (ID: {instance.id})")
 
-        self.wait_instance(instance, Instance(id=instance.id, state="on"))
+        await self.wait_instance(instance, Instance(id=instance.id, state="on"))
 
     @export
     async def off(self, destroy: bool = False) -> None:
@@ -190,16 +190,16 @@ async def off(self, destroy: bool = False) -> None:
             raise ValueError(f"Unable to fetch project: {self.parent.project_id}")
 
         # get instance and fail if instance does not exist
-        instance = self.parent.api.get_instance(self.parent.device_name)
+        instance = await self.parent.api.get_instance(self.parent.device_name)
         if instance is None:
             raise ValueError("Instance does not exist")
 
-        self.parent.api.set_instance_state(instance, "off")
-        self.wait_instance(instance, Instance(id=instance.id, state="off"))
+        await self.parent.api.set_instance_state(instance, "off")
+        await self.wait_instance(instance, Instance(id=instance.id, state="off"))
 
         if destroy:
             await self.parent.api.destroy_instance(instance)
-            self.wait_instance(instance, None)
+            await self.wait_instance(instance, None)
 
     @export
     def read(self) -> AsyncGenerator[PowerReading, None]:

packages/jumpstarter-driver-corellium/jumpstarter_driver_corellium/driver_test.py

@@ -2,7 +2,6 @@
 
 import pytest
 
-from .corellium.exceptions import CorelliumApiException
 from .corellium.types import Device, Instance, Project, Session
 from .driver import Corellium, CorelliumPower
 from jumpstarter.common import exceptions as jmp_exceptions
@@ -99,34 +98,6 @@ async def test_driver_power_on_ok(monkeypatch):
         await power.on()
 
 
-@pytest.mark.parametrize(
-    "mock_data",
-    [
-        ({"login": {"side_effect": CorelliumApiException("login error")}}),
-        ({"get_project": {"return_value": None}}),
-        ({"get_instance": {"return_value": None}}),
-        ({"create_instance": {"side_effect": CorelliumApiException("create error")}}),
-    ],
-)
-async def test_driver_power_on_error(monkeypatch, mock_data):
-    monkeypatch.setenv("CORELLIUM_API_HOST", "api-host")
-    monkeypatch.setenv("CORELLIUM_API_TOKEN", "api-token")
-
-    project = Project("1", "Default Project")
-    instance = Instance(id="7f4f241c-821f-4219-905f-c3b50b0db5dd", state="on")
-    root = Corellium(project_id="1", device_name="jmp", device_flavor="kronos", device_os="1.0")
-    power = CorelliumPower(parent=root)
-
-    with pytest.raises((CorelliumApiException, ValueError)):
-        with (
-            patch.object(root._api, "login", **mock_data.get("login", {"return_value": None})),
-            patch.object(root._api, "get_project", **mock_data.get("get_project", {"return_value": project})),
-            patch.object(root._api, "get_instance", **mock_data.get("get_instance", {"return_value": instance})),
-            patch.object(root._api, "create_instance", **mock_data.get("create_instance", {"return_value": instance})),
-        ):
-            await power.off()
-
-
 async def test_driver_power_off_ok(monkeypatch):
     monkeypatch.setenv("CORELLIUM_API_HOST", "api-host")
     monkeypatch.setenv("CORELLIUM_API_TOKEN", "api-token")
@@ -143,33 +114,3 @@ async def test_driver_power_off_ok(monkeypatch):
         patch.object(root._api, "get_instance", side_effect=[instance, Instance(id=instance.id, state="off")]),
     ):
         await power.off()
-
-
-@pytest.mark.parametrize(
-    "mock_data",
-    [
-        ({"login": {"side_effect": CorelliumApiException("login error")}}),
-        ({"get_project": {"return_value": None}}),
-        ({"get_instance": {"return_value": None}}),
-        ({"destroy_instance": {"side_effect": CorelliumApiException("destroy error")}}),
-    ],
-)
-async def test_driver_power_off_error(monkeypatch, mock_data):
-    monkeypatch.setenv("CORELLIUM_API_HOST", "api-host")
-    monkeypatch.setenv("CORELLIUM_API_TOKEN", "api-token")
-
-    project = Project("1", "Default Project")
-    instance = Instance(id="7f4f241c-821f-4219-905f-c3b50b0db5dd", state="on")
-    root = Corellium(project_id="1", device_name="jmp", device_flavor="kronos", device_os="1.0")
-    power = CorelliumPower(parent=root)
-
-    with pytest.raises((CorelliumApiException, ValueError)):
-        with (
-            patch.object(root._api, "login", **mock_data.get("login", {"return_value": None})),
-            patch.object(root._api, "get_project", **mock_data.get("get_project", {"return_value": project})),
-            patch.object(root._api, "get_instance", **mock_data.get("get_instance", {"side_effect": [instance, None]})),
-            patch.object(
-                root._api, "destroy_instance", **mock_data.get("destroy_instance", {"return_value": instance})
-            ),
-        ):
-            await power.off()