Skip to content

Commit

Permalink
fix: Fir 31676 python sdk resource manager errors when referencing en…
Browse files Browse the repository at this point in the history
…gines (#366)
  • Loading branch information
stepansergeevitch authored May 6, 2024
1 parent 77e260e commit 880b4cc
Show file tree
Hide file tree
Showing 6 changed files with 114 additions and 2 deletions.
2 changes: 1 addition & 1 deletion src/firebolt/common/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,4 +3,4 @@
DEFAULT_TIMEOUT_SECONDS: int = 60

# Running statuses in infromation schema
ENGINE_STATUS_RUNNING_LIST = ["Running", "ENGINE_STATE_RUNNING"]
ENGINE_STATUS_RUNNING_LIST = ["RUNNING", "Running", "ENGINE_STATE_RUNNING"]
12 changes: 12 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 Optional


class EngineType(Enum):
Expand Down Expand Up @@ -50,9 +51,20 @@ class EngineStatus(Enum):
FAILED = "Failed"
DELETING = "Deleting"

@classmethod
def _missing_(cls, value: object) -> Optional["EngineStatus"]:
"""Handle case-insensitive string values."""
if isinstance(value, str):
return cls(value.capitalize())
return super()._missing_(value)

def __str__(self) -> str:
return self.value

@classmethod
def from_string(cls, value: str) -> "EngineStatus":
return cls(value.capitalize())


class DatabaseOrder(Enum):
DATABASE_ORDER_UNSPECIFIED = "DATABASE_ORDER_UNSPECIFIED"
Expand Down
34 changes: 34 additions & 0 deletions tests/unit/async_db/V2/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,40 @@ async def test_connect_engine_name(
assert await connection.cursor().execute("select*") == len(python_query_data)


async def test_connect_engine_running(
httpx_mock: HTTPXMock,
db_name: str,
auth_url: str,
server: str,
auth: Auth,
account_name: str,
get_system_engine_url: str,
get_system_engine_callback: Callable,
account_id_url: str,
account_id_callback: Callable,
check_credentials_callback: Callable,
engine_name: str,
system_engine_query_url: str,
get_engine_url_callback_test_status: Callable,
) -> None:
httpx_mock.add_callback(check_credentials_callback, url=auth_url)
httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url)
httpx_mock.add_callback(account_id_callback, url=account_id_url)
httpx_mock.add_callback(
get_engine_url_callback_test_status, url=system_engine_query_url
)

async with await connect(
engine_name=engine_name,
database=db_name,
auth=auth,
account_name=account_name,
api_endpoint=server,
):
# Engine is running, no exception should be raised
pass


async def test_connect_database(
db_name: str,
auth_url: str,
Expand Down
34 changes: 34 additions & 0 deletions tests/unit/db/V2/test_connection.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,6 +141,40 @@ def test_connect_engine_name(
assert connection.cursor().execute("select*") == len(python_query_data)


def test_connect_engine_running(
httpx_mock: HTTPXMock,
db_name: str,
auth_url: str,
server: str,
auth: Auth,
account_name: str,
get_system_engine_url: str,
get_system_engine_callback: Callable,
account_id_url: str,
account_id_callback: Callable,
check_credentials_callback: Callable,
engine_name: str,
system_engine_query_url: str,
get_engine_url_callback_test_status: Callable,
) -> None:
httpx_mock.add_callback(check_credentials_callback, url=auth_url)
httpx_mock.add_callback(get_system_engine_callback, url=get_system_engine_url)
httpx_mock.add_callback(account_id_callback, url=account_id_url)
httpx_mock.add_callback(
get_engine_url_callback_test_status, url=system_engine_query_url
)

with connect(
engine_name=engine_name,
database=db_name,
auth=auth,
account_name=account_name,
api_endpoint=server,
):
# Engine is running, no exception should be raised
pass


def test_connect_database(
db_name: str,
auth_url: str,
Expand Down
7 changes: 7 additions & 0 deletions tests/unit/db_conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -303,6 +303,13 @@ def get_engine_url_not_running_callback(
return _get_engine_url_callback(engine_name, db_name, query_statistics, "Stopped")


@fixture(params=["Running", "RUNNING", "ENGINE_STATE_RUNNING"])
def get_engine_url_callback_test_status(
server: str, db_name: str, query_statistics: Dict[str, Any], request: Any
) -> Callable:
return _get_engine_url_callback(server, db_name, query_statistics, request.param)


@fixture
def get_engine_url_invalid_db_callback(
engine_name,
Expand Down
27 changes: 26 additions & 1 deletion tests/unit/service/V2/test_engine.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from typing import Callable

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

from firebolt.model.V2.database import Database
Expand Down Expand Up @@ -228,3 +228,28 @@ def test_engine_deleteting(
engine = resource_manager.engines.get_by_name(mock_engine.name)

assert engine.current_status == EngineStatus.DELETING


@mark.parametrize(
"engine_status, expected_status",
[(status.value.upper(), status) for status in EngineStatus],
)
def test_engine_new_status(
engine_status: str,
expected_status: EngineStatus,
httpx_mock: HTTPXMock,
resource_manager: ResourceManager,
instance_type_callback: Callable,
instance_type_url: str,
system_engine_no_db_query_url: str,
mock_engine: Engine,
):
mock_engine.current_status = engine_status
get_engine_callback = get_objects_from_db_callback([mock_engine])

httpx_mock.add_callback(instance_type_callback, url=instance_type_url)
httpx_mock.add_callback(get_engine_callback, url=system_engine_no_db_query_url)

engine = resource_manager.engines.get_by_name(mock_engine.name)

assert engine.current_status == expected_status

0 comments on commit 880b4cc

Please sign in to comment.