diff --git a/portal_client/auth.py b/portal_client/auth.py index 6f5e1fd..21f75df 100644 --- a/portal_client/auth.py +++ b/portal_client/auth.py @@ -26,6 +26,9 @@ # How long to wait for the user to complete the browser login, in seconds. LOGIN_TIMEOUT = 300 +# Refresh a stored access token this many seconds before it expires. +REFRESH_MARGIN = 60 + # Environment variables used to configure the OAuth client (see resolve_login_settings). CLIENT_ID_ENV = "PORTAL_BACKEND_CLIENT_ID" CLIENT_SECRET_ENV = "PORTAL_BACKEND_CLIENT_SECRET" @@ -113,11 +116,73 @@ def clear_credentials() -> bool: def get_stored_access_token() -> str | None: - """Return the access token from persisted credentials, if any.""" + """Return the access token from persisted credentials, if any. + + An expired (or nearly expired) token is first renewed with the stored refresh + token, so a single login lasts as long as the refresh token does. + """ credentials = load_credentials() - if credentials: - return credentials.get("access_token") - return None + if not credentials: + return None + expires_at = credentials.get("expires_at") + if ( + expires_at + and time.time() >= expires_at - REFRESH_MARGIN + and credentials.get("refresh_token") + and credentials.get("client_id") + ): + credentials = refresh_access_token(credentials) + return credentials.get("access_token") + + +def _credentials_from_token_response( + token_response: dict, endpoint: str, client_id: str, refresh_token: str | None +) -> dict: + """Build the persisted credentials from a token endpoint response.""" + credentials = { + "access_token": token_response["access_token"], + "token_type": token_response.get("token_type", "Bearer"), + "scope": token_response.get("scope"), + "endpoint": endpoint, + # Kept so the token can be refreshed without the user logging in again. + "client_id": client_id, + "refresh_token": token_response.get("refresh_token", refresh_token), + } + expires_in = token_response.get("expires_in") + if expires_in: + credentials["expires_at"] = time.time() + expires_in + return credentials + + +def refresh_access_token(credentials: dict) -> dict: + """Trade the stored refresh token for a new access token and persist it.""" + client_id = credentials["client_id"] + # A confidential client's secret is never stored; it comes from env or config. + client_secret = resolve_login_settings()["client_secret"] + response = requests.post( + urljoin(credentials["endpoint"], TOKEN_PATH), + data={ + "grant_type": "refresh_token", + "refresh_token": credentials["refresh_token"], + "client_id": client_id, + }, + auth=(client_id, client_secret) if client_secret else None, + timeout=30, + ) + + if not response.ok: + raise Exception( + "Refreshing the access token failed (HTTP {}): {}. Run " + "'innoactive-portal auth login' again.".format( + response.status_code, response.text + ) + ) + + refreshed = _credentials_from_token_response( + response.json(), credentials["endpoint"], client_id, credentials["refresh_token"] + ) + save_credentials(refreshed) + return refreshed def generate_pkce_pair() -> tuple[str, str]: @@ -298,17 +363,9 @@ def login( client_secret, ) - credentials = { - "access_token": token_response["access_token"], - "token_type": token_response.get("token_type", "Bearer"), - "scope": token_response.get("scope"), - "endpoint": endpoint, - } - expires_in = token_response.get("expires_in") - if expires_in: - credentials["expires_at"] = time.time() + expires_in - - save_credentials(credentials) + save_credentials( + _credentials_from_token_response(token_response, endpoint, client_id, None) + ) return token_response @@ -362,7 +419,9 @@ def status_cli(args): expires_at = credentials.get("expires_at") if expires_at: remaining = expires_at - time.time() - if remaining <= 0: + if remaining <= 0 and credentials.get("refresh_token"): + print("Access token has expired, it is refreshed on the next request.") + elif remaining <= 0: print("Access token has expired, run 'innoactive-portal auth login' again.") else: print("Access token expires in {} minutes.".format(int(remaining // 60))) diff --git a/portal_client/session_management.py b/portal_client/session_management.py index 00fd135..842c7bb 100644 --- a/portal_client/session_management.py +++ b/portal_client/session_management.py @@ -15,6 +15,10 @@ logging.getLogger("backoff").addHandler(logging.StreamHandler()) +# Session management scopes requests to an organization by this header, not a query parameter. +ORGANIZATION_HEADER = "Portal-Organization-Id" + + class SessionManagementApiClient: """ Class dealing with the session management API for Virtual Machines @@ -28,47 +32,117 @@ def __init__(self, base_url=None) -> None: @backoff.on_exception( backoff.expo, requests.exceptions.ConnectionError, max_time=60 ) - def list_vms(self, organization_id): + def _send(self, method, path, organization_id=None, **kwargs): """ - List VMs for an organization + Send an authenticated request, scoped to an organization if one is given. + Prints the error body and raises on a failed response; returns the parsed + JSON body, or None for an empty one. """ - response = requests.get( - urljoin(self.base_url, "/VirtualMachines"), - headers={"Authorization": get_bearer_authorization_header()}, - params={"organization_id": organization_id}, + headers = {"Authorization": get_bearer_authorization_header()} + if organization_id is not None: + headers[ORGANIZATION_HEADER] = str(organization_id) + + response = requests.request( + method, + urljoin(self.base_url, path), + headers=headers, timeout=30, + **kwargs, ) if not response.ok: - print(response.json()) + print(response.text) response.raise_for_status() - return response.json() + return response.json() if response.content else None + + def list_vms(self, organization_id=None): + """ + List VMs for an organization: the caller's own, or all of them for an + organization admin. Without an organization, lists the VMs that belong to + none, which needs the global list_virtual_machines permission. + """ + return self._send("GET", "/VirtualMachines", organization_id) + + def get_vm(self, vm_id, organization_id=None): + """ + Get a single VM, including its state and public IP + """ + return self._send("GET", f"/VirtualMachines/{vm_id}", organization_id) + + def create_vm( + self, + organization_id, + region, + size, + expiration, + image=None, + debug_mode=False, + ): + """ + Create a VM for the caller in an organization; it is started right away. + The expiration (HH:MM:SS) is how long until the VM is destroyed again. + Debug mode opens the debug ports and needs the can_debug permission. + """ + return self._send( + "POST", + "/VirtualMachines", + organization_id, + json={ + "region": region, + "size": size, + "image": image, + "expirationTimeSpan": expiration, + "debugModeEnabled": debug_mode, + }, + ) + + def destroy_vm(self, vm_id, organization_id=None): + """ + Destroy a VM + """ + return self._send("POST", f"/VirtualMachines/{vm_id}/Destroy", organization_id) + + def list_vm_sizes(self): + """ + List the VM sizes and the regions each is available in + """ + return self._send("GET", "/VirtualMachines/Sizes") + + def list_vm_images(self, instance=None, version=None, variant=None, gpu_type=None): + """ + List the VM images and the regions each is available in + """ + return self._send( + "GET", + "/VirtualMachines/Images", + params={ + "instance": instance, + "version": version, + "variant": variant, + "gpuType": gpu_type, + }, + ) - @backoff.on_exception( - backoff.expo, requests.exceptions.ConnectionError, max_time=60 - ) def extend_vm_expiration(self, vm_id, organization_id, timespan): """ Extend the expiration time of a VM """ - response = requests.put( - urljoin(self.base_url, f"/VirtualMachines/{vm_id}/Expiration"), - headers={"Authorization": get_bearer_authorization_header()}, - params={"organization_id": organization_id}, + return self._send( + "PUT", + f"/VirtualMachines/{vm_id}/Expiration", + organization_id, json={"time": timespan}, - timeout=30, ) - if not response.ok: - print(response.json()) - response.raise_for_status() - - return response.json() + def set_debug_mode(self, vm_id, organization_id, enabled): + """ + Enable or disable debug mode on a VM. Debug mode opens the debug ports + (RDP, WinRM, ...) in the VM's firewall; needs the can_debug permission. + """ + action = "EnableDebugMode" if enabled else "DisableDebugMode" + return self._send("POST", f"/VirtualMachines/{vm_id}/{action}", organization_id) - @backoff.on_exception( - backoff.expo, requests.exceptions.ConnectionError, max_time=60 - ) def refresh_regions(self): """ Force an immediate refresh of the cached cloud resources (subnets, @@ -76,15 +150,7 @@ def refresh_regions(self): cache expiry. Use after publishing a new VM image so it is picked up without waiting for the background refresh cycle. Admin only. """ - response = requests.post( - urljoin(self.base_url, "/Regions/refresh"), - headers={"Authorization": get_bearer_authorization_header()}, - timeout=30, - ) - - if not response.ok: - print(response.text) - response.raise_for_status() + self._send("POST", "/Regions/refresh") def list_vms_cli(args): @@ -94,6 +160,50 @@ def list_vms_cli(args): print(json.dumps(vms_response)) +def get_vm_cli(args): + """CLI wrapper for getting a single VM""" + client = SessionManagementApiClient() + print(json.dumps(client.get_vm(args.vm_id, args.org_id))) + + +def create_vm_cli(args): + """CLI wrapper for creating a VM""" + client = SessionManagementApiClient() + response = client.create_vm( + organization_id=args.org_id, + region=args.region, + size=args.size, + expiration=args.expiration, + image=args.image, + debug_mode=args.debug_mode, + ) + print(json.dumps(response)) + + +def destroy_vm_cli(args): + """CLI wrapper for destroying a VM""" + client = SessionManagementApiClient() + print(json.dumps(client.destroy_vm(args.vm_id, args.org_id))) + + +def list_vm_sizes_cli(args): + """CLI wrapper for listing VM sizes""" + client = SessionManagementApiClient() + print(json.dumps(client.list_vm_sizes())) + + +def list_vm_images_cli(args): + """CLI wrapper for listing VM images""" + client = SessionManagementApiClient() + response = client.list_vm_images( + instance=args.instance, + version=args.version, + variant=args.variant, + gpu_type=args.gpu_type, + ) + print(json.dumps(response)) + + def extend_vm_expiration_cli(args): """CLI wrapper for extending VM expiration""" client = SessionManagementApiClient() @@ -103,6 +213,18 @@ def extend_vm_expiration_cli(args): print(json.dumps(response)) +def enable_debug_mode_cli(args): + """CLI wrapper for enabling debug mode on a VM""" + client = SessionManagementApiClient() + print(json.dumps(client.set_debug_mode(args.vm_id, args.org_id, enabled=True))) + + +def disable_debug_mode_cli(args): + """CLI wrapper for disabling debug mode on a VM""" + client = SessionManagementApiClient() + print(json.dumps(client.set_debug_mode(args.vm_id, args.org_id, enabled=False))) + + def refresh_regions_cli(args): """CLI wrapper for refreshing cloud resources across all regions""" client = SessionManagementApiClient() @@ -119,10 +241,73 @@ def configure_session_management_parser(parser: argparse.ArgumentParser): # vm list command vm_list_parser = vm_parser.add_parser("list", help="List VMs for an organization") vm_list_parser.add_argument( - "--org-id", type=int, required=True, help="Organization ID to list VMs for" + "--org-id", + type=int, + help="Organization ID to list VMs for; omit to list VMs without an organization", ) vm_list_parser.set_defaults(func=list_vms_cli) + # vm get command + vm_get_parser = vm_parser.add_parser("get", help="Get a single VM") + vm_get_parser.add_argument("vm_id", help="ID of the VM") + vm_get_parser.add_argument( + "--org-id", type=int, help="Organization ID the VM belongs to" + ) + vm_get_parser.set_defaults(func=get_vm_cli) + + # vm create command + vm_create_parser = vm_parser.add_parser( + "create", help="Create a VM for yourself in an organization, and start it" + ) + vm_create_parser.add_argument( + "--org-id", type=int, required=True, help="Organization ID to create it in" + ) + vm_create_parser.add_argument( + "--region", required=True, help="Cloud region, e.g. eu-central-1" + ) + vm_create_parser.add_argument( + "--size", required=True, help="VM size name, see 'vms sizes'" + ) + vm_create_parser.add_argument( + "--expiration", + required=True, + help="Time until the VM is destroyed again, in format HH:MM:SS", + ) + vm_create_parser.add_argument( + "--image", + help="Image, e.g. dev/latest, see 'vms images'; defaults to the region's default", + ) + vm_create_parser.add_argument( + "--debug-mode", + action="store_true", + help="Create it with debug mode on, opening its debug ports (RDP, WinRM)", + ) + vm_create_parser.set_defaults(func=create_vm_cli) + + # vm destroy command + vm_destroy_parser = vm_parser.add_parser("destroy", help="Destroy a VM") + vm_destroy_parser.add_argument("vm_id", help="ID of the VM") + vm_destroy_parser.add_argument( + "--org-id", type=int, help="Organization ID the VM belongs to" + ) + vm_destroy_parser.set_defaults(func=destroy_vm_cli) + + # vm sizes command + vm_sizes_parser = vm_parser.add_parser( + "sizes", help="List VM sizes and the regions they are available in" + ) + vm_sizes_parser.set_defaults(func=list_vm_sizes_cli) + + # vm images command + vm_images_parser = vm_parser.add_parser( + "images", help="List VM images and the regions they are available in" + ) + vm_images_parser.add_argument("--instance", help="Filter by instance, e.g. dev") + vm_images_parser.add_argument("--version", help="Filter by version, e.g. 1.0.0") + vm_images_parser.add_argument("--variant", help="Filter by variant") + vm_images_parser.add_argument("--gpu-type", help="Filter by GPU type, e.g. l4") + vm_images_parser.set_defaults(func=list_vm_images_cli) + # vm extend-expiration command vm_extend_parser = vm_parser.add_parser( "extend-expiration", help="Extend the expiration time of a VM" @@ -136,6 +321,26 @@ def configure_session_management_parser(parser: argparse.ArgumentParser): ) vm_extend_parser.set_defaults(func=extend_vm_expiration_cli) + # vm enable-debug-mode / disable-debug-mode commands + for name, func, help_text in ( + ( + "enable-debug-mode", + enable_debug_mode_cli, + "Enable debug mode on a VM, opening its debug ports (RDP, WinRM)", + ), + ( + "disable-debug-mode", + disable_debug_mode_cli, + "Disable debug mode on a VM, closing its debug ports again", + ), + ): + debug_parser = vm_parser.add_parser(name, help=help_text) + debug_parser.add_argument("vm_id", help="ID of the VM") + debug_parser.add_argument( + "--org-id", type=int, help="Organization ID the VM belongs to" + ) + debug_parser.set_defaults(func=func) + return vm_parser diff --git a/readme.md b/readme.md index 8080d5d..bb429a4 100644 --- a/readme.md +++ b/readme.md @@ -87,7 +87,10 @@ innoactive-portal auth login ``` The obtained token is stored in `~/.config/innoactive-portal/credentials.json` (honoring -`XDG_CONFIG_HOME`). To check your current status or log out again, use: +`XDG_CONFIG_HOME`). When the server issues a refresh token, it is stored alongside, and an +expired access token is renewed with it on the next request, so a login lasts as long as +the refresh token does instead of the access token's lifetime. To check your current +status or log out again, use: ```sh innoactive-portal auth status @@ -145,6 +148,35 @@ innoactive-portal applications v2 upload-build \ You can run `innoactive-portal applications v2 upload-build --help` to get more information on available parameters. +### Managing virtual machines + +Virtual machines live in session management, set `PORTAL_SESSION_MANAGEMENT_ENDPOINT` for +an instance other than production (see [Configuration](#configuration)). A VM that belongs +to an organization is only found when you pass that organization's ID with `--org-id`. +Without org admin rights, you only see your own VMs. + +```sh +# List your VMs in organization 1 (omit --org-id for the VMs without an organization) +innoactive-portal vms list --org-id 1 + +# Find a size and an image for a region +innoactive-portal vms sizes +innoactive-portal vms images --instance dev --gpu-type t4 + +# Create a VM for yourself (it starts right away) that is destroyed again after an hour, +# with debug mode on so its debug ports (RDP, WinRM) are open +innoactive-portal vms create --org-id 1 --region eu-central-1 --size t4.small --image dev/latest --expiration 01:00:00 --debug-mode + +# Poll its state and public IP, toggle debug mode, destroy it +innoactive-portal vms get --org-id 1 +innoactive-portal vms disable-debug-mode --org-id 1 +innoactive-portal vms enable-debug-mode --org-id 1 +innoactive-portal vms destroy --org-id 1 +``` + +Debug mode requires the `can_debug` permission. Session management has no endpoint to +start a stopped VM; a VM starts when it is created or when a session is placed on it. + ### Refreshing region resources Forces an immediate refresh of the cached cloud resources (subnets, gateways, and VM images) across all regions in session management, bypassing the configured cache expiry. Useful after publishing a new VM image so it is picked up without waiting for the background refresh cycle. Requires admin access. diff --git a/tests/test_auth.py b/tests/test_auth.py index 5580598..5ced3fe 100644 --- a/tests/test_auth.py +++ b/tests/test_auth.py @@ -1,4 +1,5 @@ import hashlib +import time from base64 import urlsafe_b64encode from unittest.mock import patch from urllib.parse import parse_qs, urlparse @@ -265,6 +266,97 @@ def test_login_raises_on_authorization_error(self, isolated_config): ) +class TestTokenRefresh: + def test_login_keeps_refresh_token_and_client(self, requests_mock, isolated_config): + requests_mock.post( + "https://api.innoactive.io/oauth/token/", + json={"access_token": "a1", "refresh_token": "r1", "expires_in": 3600}, + ) + + with patch( + "portal_client.auth._wait_for_callback", + return_value={"code": "c", "state": "s", "error": None}, + ), patch("portal_client.auth.secrets.token_urlsafe", side_effect=["v", "s"]): + auth.login( + client_id="my-client", + endpoint="https://api.innoactive.io", + open_browser=False, + ) + + stored = auth.load_credentials() + assert stored["refresh_token"] == "r1" + assert stored["client_id"] == "my-client" + + def test_valid_token_is_not_refreshed(self, requests_mock, isolated_config): + auth.save_credentials( + { + "access_token": "a1", + "refresh_token": "r1", + "client_id": "my-client", + "endpoint": "https://api.innoactive.io", + "expires_at": time.time() + 3600, + } + ) + + assert auth.get_stored_access_token() == "a1" + assert not requests_mock.called + + def test_expired_token_is_refreshed_and_persisted( + self, requests_mock, isolated_config, monkeypatch + ): + monkeypatch.delenv("PORTAL_BACKEND_CLIENT_SECRET", raising=False) + requests_mock.post( + "https://api.innoactive.io/oauth/token/", + json={"access_token": "a2", "refresh_token": "r2", "expires_in": 3600}, + ) + auth.save_credentials( + { + "access_token": "a1", + "refresh_token": "r1", + "client_id": "my-client", + "endpoint": "https://api.innoactive.io", + "expires_at": time.time() - 1, + } + ) + + assert auth.get_stored_access_token() == "a2" + body = parse_qs(requests_mock.last_request.text) + assert body == { + "grant_type": ["refresh_token"], + "refresh_token": ["r1"], + "client_id": ["my-client"], + } + # The rotated refresh token replaces the used one. + assert auth.load_credentials()["refresh_token"] == "r2" + + def test_expired_token_without_refresh_token_is_returned_as_is( + self, requests_mock, isolated_config + ): + auth.save_credentials({"access_token": "a1", "expires_at": time.time() - 1}) + + assert auth.get_stored_access_token() == "a1" + assert not requests_mock.called + + def test_failed_refresh_asks_for_login(self, requests_mock, isolated_config): + requests_mock.post( + "https://api.innoactive.io/oauth/token/", + json={"error": "invalid_grant"}, + status_code=400, + ) + auth.save_credentials( + { + "access_token": "a1", + "refresh_token": "r1", + "client_id": "my-client", + "endpoint": "https://api.innoactive.io", + "expires_at": time.time() - 1, + } + ) + + with pytest.raises(Exception, match="auth login"): + auth.get_stored_access_token() + + class TestBearerHeaderFallback: def test_env_var_takes_precedence(self, isolated_config, monkeypatch): from portal_client import utils diff --git a/tests/test_session_management.py b/tests/test_session_management.py index 3f01ceb..1a27af4 100644 --- a/tests/test_session_management.py +++ b/tests/test_session_management.py @@ -1,13 +1,15 @@ import json from io import StringIO from unittest.mock import patch -from urllib.parse import parse_qs, urlparse import pytest import requests from portal_client.session_management import ( SessionManagementApiClient, + create_vm_cli, + disable_debug_mode_cli, + enable_debug_mode_cli, extend_vm_expiration_cli, list_vms_cli, refresh_regions_cli, @@ -37,15 +39,28 @@ def test_list_vms_success(self, requests_mock): result = client.list_vms(organization_id=123) assert result == expected_response - # Check the request URL contains the correct parameters - request_url = requests_mock.last_request.url - parsed_url = urlparse(request_url) - query_params = parse_qs(parsed_url.query) - assert query_params["organization_id"][0] == "123" + # Session management only reads the organization from this header + assert requests_mock.last_request.headers["Portal-Organization-Id"] == "123" + assert requests_mock.last_request.qs == {} assert ( requests_mock.last_request.headers["Authorization"] == "Bearer test-token" ) + def test_list_vms_without_organization_sends_no_header(self, requests_mock): + requests_mock.get( + "https://session-management.innoactive.io/VirtualMachines", + json={"items": []}, + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + client.list_vms() + + assert "Portal-Organization-Id" not in requests_mock.last_request.headers + def test_extend_vm_expiration_success(self, requests_mock): # Mock the API response expected_response = {"message": "VM expiration extended successfully"} @@ -65,12 +80,9 @@ def test_extend_vm_expiration_success(self, requests_mock): ) assert result == expected_response - # Check the request URL contains the correct parameters - request_url = requests_mock.last_request.url - parsed_url = urlparse(request_url) - query_params = parse_qs(parsed_url.query) request_body = requests_mock.last_request.json() - assert query_params["organization_id"][0] == "456" + assert requests_mock.last_request.headers["Portal-Organization-Id"] == "456" + assert requests_mock.last_request.qs == {} assert request_body["time"] == "01:00:00" assert ( requests_mock.last_request.headers["Authorization"] == "Bearer test-token" @@ -91,6 +103,133 @@ def test_list_vms_with_custom_base_url(self, requests_mock): assert result == expected_response + @pytest.mark.parametrize( + "enabled, action", [(True, "EnableDebugMode"), (False, "DisableDebugMode")] + ) + def test_set_debug_mode(self, requests_mock, enabled, action): + expected_response = {"id": "vm-123", "debugModeEnabled": enabled} + requests_mock.post( + f"https://session-management.innoactive.io/VirtualMachines/vm-123/{action}", + json=expected_response, + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + result = client.set_debug_mode("vm-123", 456, enabled=enabled) + + assert result == expected_response + assert requests_mock.last_request.headers["Portal-Organization-Id"] == "456" + + def test_set_debug_mode_error_response(self, requests_mock): + requests_mock.post( + "https://session-management.innoactive.io/VirtualMachines/vm-123/EnableDebugMode", + text="Forbidden", + status_code=403, + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + with pytest.raises(requests.HTTPError): + client.set_debug_mode("vm-123", 456, enabled=True) + + def test_get_vm(self, requests_mock): + requests_mock.get( + "https://session-management.innoactive.io/VirtualMachines/vm-123", + json={"id": "vm-123", "publicIp": "1.2.3.4"}, + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + result = client.get_vm("vm-123", 456) + + assert result == {"id": "vm-123", "publicIp": "1.2.3.4"} + assert requests_mock.last_request.headers["Portal-Organization-Id"] == "456" + + def test_create_vm(self, requests_mock): + requests_mock.post( + "https://session-management.innoactive.io/VirtualMachines", + json={"id": "vm-123", "state": "Created"}, + status_code=201, + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + result = client.create_vm( + organization_id=456, + region="eu-central-1", + size="l4-small", + expiration="01:00:00", + image="dev/latest", + debug_mode=True, + ) + + assert result == {"id": "vm-123", "state": "Created"} + assert requests_mock.last_request.headers["Portal-Organization-Id"] == "456" + assert requests_mock.last_request.json() == { + "region": "eu-central-1", + "size": "l4-small", + "image": "dev/latest", + "expirationTimeSpan": "01:00:00", + "debugModeEnabled": True, + } + + def test_destroy_vm(self, requests_mock): + requests_mock.post( + "https://session-management.innoactive.io/VirtualMachines/vm-123/Destroy", + json={"id": "vm-123", "state": "Started"}, + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + client.destroy_vm("vm-123", 456) + + assert requests_mock.last_request.method == "POST" + assert requests_mock.last_request.headers["Portal-Organization-Id"] == "456" + + def test_list_vm_sizes(self, requests_mock): + expected_response = {"vmSizes": [{"name": "l4-small", "isEnabled": True}]} + requests_mock.get( + "https://session-management.innoactive.io/VirtualMachines/Sizes", + json=expected_response, + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + assert client.list_vm_sizes() == expected_response + + def test_list_vm_images_passes_only_given_filters(self, requests_mock): + requests_mock.get( + "https://session-management.innoactive.io/VirtualMachines/Images", + json=[], + ) + + client = SessionManagementApiClient() + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + client.list_vm_images(instance="dev", gpu_type="l4") + + assert requests_mock.last_request.qs == {"instance": ["dev"], "gputype": ["l4"]} + def test_refresh_regions_success(self, requests_mock): # The endpoint returns 200 OK with no content body requests_mock.post( @@ -221,3 +360,55 @@ class MockArgs: output = mock_stdout.getvalue().strip() assert "Successfully triggered a refresh" in output assert requests_mock.last_request.method == "POST" + + @pytest.mark.parametrize( + "cli, action, enabled", + [ + (enable_debug_mode_cli, "EnableDebugMode", True), + (disable_debug_mode_cli, "DisableDebugMode", False), + ], + ) + def test_debug_mode_cli(self, requests_mock, cli, action, enabled): + expected_response = {"id": "vm-123", "debugModeEnabled": enabled} + requests_mock.post( + f"https://session-management.innoactive.io/VirtualMachines/vm-123/{action}", + json=expected_response, + ) + + class MockArgs: + vm_id = "vm-123" + org_id = 456 + + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + with patch("sys.stdout", new_callable=StringIO) as mock_stdout: + cli(MockArgs()) + + assert json.loads(mock_stdout.getvalue().strip()) == expected_response + + def test_create_vm_cli(self, requests_mock): + requests_mock.post( + "https://session-management.innoactive.io/VirtualMachines", + json={"id": "vm-123"}, + status_code=201, + ) + + class MockArgs: + org_id = 456 + region = "eu-central-1" + size = "l4-small" + expiration = "01:00:00" + image = None + debug_mode = False + + with patch( + "portal_client.session_management.get_bearer_authorization_header", + return_value="Bearer test-token", + ): + with patch("sys.stdout", new_callable=StringIO) as mock_stdout: + create_vm_cli(MockArgs()) + + assert json.loads(mock_stdout.getvalue().strip()) == {"id": "vm-123"} + assert requests_mock.last_request.json()["image"] is None