fix: Readd get_many for old engines (#375)

Author: ptiurin

setup.cfg

@@ -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

src/firebolt/service/V1/engine.py

@@ -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__)
 
@@ -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"]
+        ]

tests/integration/resource_manager/V1/test_engine.py

@@ -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)

tests/unit/service/V1/conftest.py

@@ -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(

tests/unit/service/V1/test_engine.py

@@ -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(
@@ -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