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
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,7 @@ ciso8601 =
dev =
allure-pytest==2.*
devtools==0.7.0
mypy==1.*
mypy==1.*,<1.10.0
pre-commit==2.15.0
pyfakefs>=4.5.3
pytest==7.2.0
Expand Down
53 changes: 52 additions & 1 deletion src/firebolt/service/V1/engine.py
Original file line number Diff line number Diff line change
@@ -1,13 +1,16 @@
from logging import getLogger
from typing import List
from typing import List, Optional, Union

from firebolt.model.V1.engine import Engine
from firebolt.service.V1.base import BaseService
from firebolt.service.V1.types import EngineOrder
from firebolt.utils.urls import (
ACCOUNT_ENGINE_ID_BY_NAME_URL,
ACCOUNT_ENGINE_URL,
ACCOUNT_LIST_ENGINES_URL,
ENGINES_BY_IDS_URL,
)
from firebolt.utils.util import prune_dict

logger = getLogger(__name__)

Expand Down Expand Up @@ -47,3 +50,51 @@ def get_by_name(self, name: str) -> Engine:
)
engine_id = response.json()["engine_id"]["engine_id"]
return self.get(id_=engine_id)

def get_many(
self,
name_contains: Optional[str] = None,
current_status_eq: Optional[str] = None,
current_status_not_eq: Optional[str] = None,
region_eq: Optional[str] = None,
order_by: Optional[Union[str, EngineOrder]] = None,
) -> List[Engine]:
"""
Get a list of engines on Firebolt.

Args:
name_contains: Filter for engines with a name containing this substring
current_status_eq: Filter for engines with this status
current_status_not_eq: Filter for engines that do not have this status
region_eq: Filter for engines by region
order_by: Method by which to order the results. See [EngineOrder]

Returns:
A list of engines matching the filters
"""

if isinstance(order_by, str):
order_by = EngineOrder[order_by].name

if region_eq is not None:
region_eq = self.resource_manager.regions.get_by_name(
name=region_eq
).key.region_id

response = self.client.get(
url=ACCOUNT_LIST_ENGINES_URL.format(account_id=self.account_id),
params=prune_dict(
{
"page.first": 5000,
"filter.name_contains": name_contains,
"filter.current_status_eq": current_status_eq,
"filter.current_status_not_eq": current_status_not_eq,
"filter.compute_region_id_region_id_eq": region_eq,
"order_by": order_by,
}
),
)
return [
Engine.parse_obj_with_service(obj=e["node"], engine_service=self)
for e in response.json()["edges"]
]
6 changes: 6 additions & 0 deletions tests/integration/resource_manager/V1/test_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,12 @@ def test_get_engine(resource_manager: ResourceManager, engine_name: str):
)


def test_get_many(resource_manager: ResourceManager, engine_name: str):
engines = resource_manager.engines.get_many(name_contains=engine_name)
# >= 1 because we may have a stopped engine with the same name + _stopped
assert len(engines) >= 1


@mark.skip(reason="Interferes with other tests")
def test_engine_stop_start(resource_manager: ResourceManager, engine_name: str):
engine = resource_manager.engines.get_by_name(engine_name)
Expand Down
14 changes: 14 additions & 0 deletions tests/unit/service/V1/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -260,6 +260,20 @@ def do_mock(
return do_mock


@fixture
def many_engines_callback(mock_engine: Engine) -> Callable:
def do_mock(
request: httpx.Request = None,
**kwargs,
) -> Response:
return Response(
status_code=httpx.codes.OK,
json={"edges": [{"node": mock_engine.model_dict()}]},
)

return do_mock


@fixture
def account_engine_url(server: str, account_id, mock_engine) -> str:
return f"https://{server}" + ACCOUNT_ENGINE_URL.format(
Expand Down
56 changes: 55 additions & 1 deletion tests/unit/service/V1/test_engine.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,6 @@ def test_engine_start_stop(
httpx_mock.add_callback(auth_callback, url=auth_url)
httpx_mock.add_callback(provider_callback, url=provider_url)
httpx_mock.add_callback(account_id_callback, url=account_id_url)
httpx_mock.add_callback(auth_callback, url=auth_url)

httpx_mock.add_callback(engine_callback, url=f"{account_engine_url}", method="GET")
httpx_mock.add_callback(
Expand All @@ -44,3 +43,58 @@ def test_engine_start_stop(
engine = mock_engine.stop(wait_for_stop=False)

assert engine.name == mock_engine.name


def test_engine_get_many(
httpx_mock: HTTPXMock,
settings: Settings,
mock_engine: Engine,
auth_callback: Callable,
auth_url: str,
engine_url: str,
many_engines_callback: Callable,
provider_callback: Callable,
provider_url: str,
account_id_callback: Callable,
account_id_url: Pattern,
):

httpx_mock.add_callback(auth_callback, url=auth_url)
httpx_mock.add_callback(provider_callback, url=provider_url)
httpx_mock.add_callback(account_id_callback, url=account_id_url)
filters = "?page.first=5000&filter.name_contains=test"
httpx_mock.add_callback(many_engines_callback, url=engine_url + filters)

manager = ResourceManager(settings=settings)
engines = manager.engines.get_many(name_contains="test")
assert len(engines) == 1
assert engines[0].name == mock_engine.name


def test_engine_get_many_with_region(
httpx_mock: HTTPXMock,
settings: Settings,
mock_engine: Engine,
auth_callback: Callable,
auth_url: str,
engine_url: str,
many_engines_callback: Callable,
provider_callback: Callable,
provider_url: str,
account_id_callback: Callable,
account_id_url: Pattern,
region_callback: Callable,
region_url: str,
):

httpx_mock.add_callback(auth_callback, url=auth_url)
httpx_mock.add_callback(provider_callback, url=provider_url)
httpx_mock.add_callback(account_id_callback, url=account_id_url)
httpx_mock.add_callback(region_callback, url=region_url)
filters = "?page.first=5000&filter.name_contains=test&filter.compute_region_id_region_id_eq=mock_region_id_1"
httpx_mock.add_callback(many_engines_callback, url=engine_url + filters)

manager = ResourceManager(settings=settings)
engines = manager.engines.get_many(name_contains="test", region_eq="mock_region_1")
assert len(engines) == 1
assert engines[0].name == mock_engine.name