diff --git a/src/firebolt/model/V2/engine.py b/src/firebolt/model/V2/engine.py index 58e4c433d0b..38f567f1d99 100644 --- a/src/firebolt/model/V2/engine.py +++ b/src/firebolt/model/V2/engine.py @@ -56,10 +56,10 @@ class Engine(FireboltBaseModel): def __post_init__(self) -> None: # Specs are just strings for accounts v2 - if isinstance(self.spec, str) and self.spec: + if isinstance(self.spec, str): # Resolve engine specification self.spec = InstanceType(self.spec) - if isinstance(self.current_status, str) and self.current_status: + if isinstance(self.current_status, str): # Resolve engine status self.current_status = EngineStatus(self.current_status) diff --git a/src/firebolt/model/V2/instance_type.py b/src/firebolt/model/V2/instance_type.py index 55973aa3904..f26aeedfceb 100644 --- a/src/firebolt/model/V2/instance_type.py +++ b/src/firebolt/model/V2/instance_type.py @@ -1,4 +1,5 @@ from enum import Enum +from typing import Any class InstanceType(Enum): @@ -6,3 +7,8 @@ class InstanceType(Enum): M = "M" L = "L" XL = "XL" + UNKNOWN = "UNKNOWN" # instance type could not be determined + + @classmethod + def _missing_(cls, value: Any) -> "InstanceType": + return cls.UNKNOWN diff --git a/src/firebolt/service/V2/types.py b/src/firebolt/service/V2/types.py index 96163e05521..edaacd1632a 100644 --- a/src/firebolt/service/V2/types.py +++ b/src/firebolt/service/V2/types.py @@ -1,4 +1,5 @@ from enum import Enum +from typing import Any class EngineStatus(Enum): @@ -15,6 +16,13 @@ class EngineStatus(Enum): REPAIRING = "REPAIRING" FAILED = "FAILED" DELETING = "DELETING" + RESIZING = "RESIZING" + DRAINING = "DRAINING" + UNKNOWN = "UNKNOWN" # status could not be determined + + @classmethod + def _missing_(cls, value: Any) -> "EngineStatus": + return cls.UNKNOWN def __str__(self) -> str: return self.value diff --git a/tests/unit/service/test_engine.py b/tests/unit/service/test_engine.py index eb29783d0c8..f60db590bf7 100644 --- a/tests/unit/service/test_engine.py +++ b/tests/unit/service/test_engine.py @@ -1,4 +1,4 @@ -from typing import Callable +from typing import Callable, Union from httpx import Request from pytest import mark, raises @@ -6,6 +6,7 @@ from firebolt.model.V2.database import Database from firebolt.model.V2.engine import Engine, EngineStatus +from firebolt.model.V2.instance_type import InstanceType from firebolt.service.manager import ResourceManager from firebolt.utils.exception import EngineNotFoundError from tests.unit.response import Response @@ -219,3 +220,59 @@ def test_engine_new_status( engine = resource_manager.engines.get_by_name(mock_engine.name) assert engine.current_status == expected_status + + +@mark.parametrize( + "spec_input, expected_spec, status_input, expected_status", + [ + # Test all valid InstanceType values + ("S", InstanceType.S, "RUNNING", EngineStatus.RUNNING), + ("M", InstanceType.M, "STOPPED", EngineStatus.STOPPED), + ("L", InstanceType.L, "STARTING", EngineStatus.STARTING), + ("XL", InstanceType.XL, "STOPPING", EngineStatus.STOPPING), + # Test InstanceType enum values directly + (InstanceType.S, InstanceType.S, "FAILED", EngineStatus.FAILED), + (InstanceType.M, InstanceType.M, "REPAIRING", EngineStatus.REPAIRING), + # Test unknown/invalid values that should default to UNKNOWN + ("INVALID_TYPE", InstanceType.UNKNOWN, "INVALID_STATUS", EngineStatus.UNKNOWN), + ("XXL", InstanceType.UNKNOWN, "WEIRD_STATE", EngineStatus.UNKNOWN), + # Test empty strings that should default to UNKNOWN + ("", InstanceType.UNKNOWN, "", EngineStatus.UNKNOWN), + # Test all valid EngineStatus values with M instance type + ("M", InstanceType.M, "STARTED", EngineStatus.STARTED), + ("M", InstanceType.M, "DROPPING", EngineStatus.DROPPING), + ("M", InstanceType.M, "DELETING", EngineStatus.DELETING), + ("M", InstanceType.M, "RESIZING", EngineStatus.RESIZING), + ("M", InstanceType.M, "DRAINING", EngineStatus.DRAINING), + ], +) +def test_engine_instantiation_with_different_configurations( + spec_input: Union[str, InstanceType], + expected_spec: InstanceType, + status_input: str, + expected_status: EngineStatus, +) -> None: + """ + Test that Engine model correctly handles different instance types and statuses, + including unknown values and empty strings that should default to UNKNOWN. + """ + engine = Engine( + name="test_engine", + region="us-east-1", + spec=spec_input, + scale=2, + current_status=status_input, + version="1.0", + endpoint="https://test.endpoint.com", + warmup="", + auto_stop=3600, + type="general_purpose", + _database_name="test_db", + _service=None, + ) + + assert engine.spec == expected_spec + assert engine.current_status == expected_status + assert engine.name == "test_engine" + assert engine.region == "us-east-1" + assert engine.scale == 2 diff --git a/tests/unit/service/test_instance_type.py b/tests/unit/service/test_instance_type.py index d7a37e513a2..ab01cd45262 100644 --- a/tests/unit/service/test_instance_type.py +++ b/tests/unit/service/test_instance_type.py @@ -8,9 +8,10 @@ def test_instance_types(resource_manager: ResourceManager): InstanceType.M, InstanceType.L, InstanceType.XL, + InstanceType.UNKNOWN, ] assert resource_manager.instance_types.get("XL") == InstanceType.XL - assert resource_manager.instance_types.get("XS") is None + assert resource_manager.instance_types.get("XS") == InstanceType.UNKNOWN assert resource_manager.instance_types.cheapest_instance() == InstanceType.S