Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions src/firebolt/model/V2/engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
6 changes: 6 additions & 0 deletions src/firebolt/model/V2/instance_type.py
Original file line number Diff line number Diff line change
@@ -1,8 +1,14 @@
from enum import Enum
from typing import Any


class InstanceType(Enum):
S = "S"
M = "M"
L = "L"
XL = "XL"
UNKNOWN = "UNKNOWN" # instance type could not be determined
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are not the only supported Types, we have others in certain cases, SDKs should not break due to this. Please treat them as a string.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This is exactly what this PR addresses - not breaking due to unexpected values.


@classmethod
def _missing_(cls, value: Any) -> "InstanceType":
return cls.UNKNOWN
8 changes: 8 additions & 0 deletions src/firebolt/service/V2/types.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from enum import Enum
from typing import Any


class EngineStatus(Enum):
Expand All @@ -15,6 +16,13 @@ class EngineStatus(Enum):
REPAIRING = "REPAIRING"
FAILED = "FAILED"
DELETING = "DELETING"
RESIZING = "RESIZING"
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

REPAIRING = "REPAIRING", STARTED = "STARTED", FAILED = "FAILED" and DROPPING = "DROPPING" does not exist in V2

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These are not set by the user, but are more informational. However, it will be cleaner to remove unused settings - agreed.

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
59 changes: 58 additions & 1 deletion tests/unit/service/test_engine.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
from typing import Callable
from typing import Callable, Union

from httpx import Request
from pytest import mark, raises
from pytest_httpx import HTTPXMock

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
Expand Down Expand Up @@ -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),
Comment on lines +237 to +238
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These 2 do not exist today, but they can. We also have type families, not sure whether SDKs care, though.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can't really support what does not exist today. Once they're added on the backend we can add them here as well.
Filed https://packboard.atlassian.net/browse/FIR-50661 for type families.

# 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
3 changes: 2 additions & 1 deletion tests/unit/service/test_instance_type.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Loading