fixed tests in ck

sanderegg · Dec 2, 2024 · d7d6c82 · d7d6c82
1 parent 73b1f4c
commit d7d6c82
Show file tree

Hide file tree

Showing 5 changed files with 22 additions and 20 deletions.
diff --git a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/dask.py b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/dask.py
@@ -3,7 +3,7 @@
 from typing import Any, Final
 
 import distributed
-from models_library.clusters import InternalClusterAuthentication, TLSAuthentication
+from models_library.clusters import ClusterAuthentication, TLSAuthentication
 from pydantic import AnyUrl
 
 _logger = logging.getLogger(__name__)
@@ -21,9 +21,7 @@ async def _wrap_client_async_routine(
 _CONNECTION_TIMEOUT: Final[str] = "5"
 
 
-async def ping_scheduler(
-    url: AnyUrl, authentication: InternalClusterAuthentication
-) -> bool:
+async def ping_scheduler(url: AnyUrl, authentication: ClusterAuthentication) -> bool:
     try:
         security = distributed.Security()
         if isinstance(authentication, TLSAuthentication):
@@ -47,9 +45,7 @@ async def ping_scheduler(
     return False
 
 
-async def is_scheduler_busy(
-    url: AnyUrl, authentication: InternalClusterAuthentication
-) -> bool:
+async def is_scheduler_busy(url: AnyUrl, authentication: ClusterAuthentication) -> bool:
     security = distributed.Security()
     if isinstance(authentication, TLSAuthentication):
         security = distributed.Security(

diff --git a/services/clusters-keeper/src/simcore_service_clusters_keeper/utils/clusters.py b/services/clusters-keeper/src/simcore_service_clusters_keeper/utils/clusters.py
@@ -14,7 +14,7 @@
     ClusterState,
     OnDemandCluster,
 )
-from models_library.clusters import InternalClusterAuthentication, TLSAuthentication
+from models_library.clusters import ClusterAuthentication, TLSAuthentication
 from models_library.users import UserID
 from models_library.wallets import WalletID
 from types_aiobotocore_ec2.literals import InstanceStateNameType
@@ -190,7 +190,7 @@ def create_cluster_from_ec2_instance(
     wallet_id: WalletID | None,
     *,
     dask_scheduler_ready: bool,
-    cluster_auth: InternalClusterAuthentication,
+    cluster_auth: ClusterAuthentication,
     max_cluster_start_time: datetime.timedelta,
 ) -> OnDemandCluster:
     return OnDemandCluster(

diff --git a/services/clusters-keeper/src/simcore_service_clusters_keeper/utils/dask.py b/services/clusters-keeper/src/simcore_service_clusters_keeper/utils/dask.py
@@ -1,6 +1,6 @@
 from aws_library.ec2 import EC2InstanceData
 from fastapi import FastAPI
-from models_library.clusters import InternalClusterAuthentication
+from models_library.clusters import ClusterAuthentication
 from pydantic import AnyUrl, TypeAdapter
 
 from ..core.settings import get_application_settings
@@ -13,7 +13,7 @@ def get_scheduler_url(ec2_instance: EC2InstanceData) -> AnyUrl:
     return url
 
 
-def get_scheduler_auth(app: FastAPI) -> InternalClusterAuthentication:
+def get_scheduler_auth(app: FastAPI) -> ClusterAuthentication:
     return get_application_settings(
         app
     ).CLUSTERS_KEEPER_COMPUTATIONAL_BACKEND_DEFAULT_CLUSTER_AUTH
diff --git a/services/clusters-keeper/tests/unit/test_modules_dask.py b/services/clusters-keeper/tests/unit/test_modules_dask.py
@@ -8,7 +8,7 @@
 from distributed import SpecCluster
 from faker import Faker
 from models_library.clusters import (
-    InternalClusterAuthentication,
+    ClusterAuthentication,
     NoAuthentication,
     TLSAuthentication,
 )
@@ -34,11 +34,13 @@
     "authentication", _authentication_types, ids=lambda p: f"authentication-{p.type}"
 )
 async def test_ping_scheduler_non_existing_scheduler(
-    faker: Faker, authentication: InternalClusterAuthentication
+    faker: Faker, authentication: ClusterAuthentication
 ):
     assert (
         await ping_scheduler(
-            TypeAdapter(AnyUrl).validate_python(f"tcp://{faker.ipv4()}:{faker.port_number()}"),
+            TypeAdapter(AnyUrl).validate_python(
+                f"tcp://{faker.ipv4()}:{faker.port_number()}"
+            ),
             authentication,
         )
         is False
@@ -48,7 +50,9 @@ async def test_ping_scheduler_non_existing_scheduler(
 async def test_ping_scheduler(dask_spec_local_cluster: SpecCluster):
     assert (
         await ping_scheduler(
-            TypeAdapter(AnyUrl).validate_python(dask_spec_local_cluster.scheduler_address),
+            TypeAdapter(AnyUrl).validate_python(
+                dask_spec_local_cluster.scheduler_address
+            ),
             NoAuthentication(),
         )
         is True
@@ -71,7 +75,9 @@ async def test_is_scheduler_busy(
     dask_spec_cluster_client: distributed.Client,
 ):
     # nothing runs right now
-    scheduler_address = TypeAdapter(AnyUrl).validate_python(dask_spec_local_cluster.scheduler_address)
+    scheduler_address = TypeAdapter(AnyUrl).validate_python(
+        dask_spec_local_cluster.scheduler_address
+    )
     assert await is_scheduler_busy(scheduler_address, NoAuthentication()) is False
     _SLEEP_TIME = 5
 

diff --git a/services/clusters-keeper/tests/unit/test_utils_clusters.py b/services/clusters-keeper/tests/unit/test_utils_clusters.py
@@ -19,7 +19,7 @@
 from faker import Faker
 from models_library.api_schemas_clusters_keeper.clusters import ClusterState
 from models_library.clusters import (
-    InternalClusterAuthentication,
+    ClusterAuthentication,
     NoAuthentication,
     TLSAuthentication,
 )
@@ -55,15 +55,15 @@ def ec2_boot_specs(app_settings: ApplicationSettings) -> EC2InstanceBootSpecific
 @pytest.fixture(params=[TLSAuthentication, NoAuthentication])
 def backend_cluster_auth(
     request: pytest.FixtureRequest,
-) -> InternalClusterAuthentication:
+) -> ClusterAuthentication:
     return request.param
 
 
 @pytest.fixture
 def app_environment(
     app_environment: EnvVarsDict,
     monkeypatch: pytest.MonkeyPatch,
-    backend_cluster_auth: InternalClusterAuthentication,
+    backend_cluster_auth: ClusterAuthentication,
 ) -> EnvVarsDict:
     return app_environment | setenvs_from_dict(
         monkeypatch,
@@ -295,7 +295,7 @@ def test_create_cluster_from_ec2_instance(
     faker: Faker,
     ec2_state: InstanceStateNameType,
     expected_cluster_state: ClusterState,
-    authentication: InternalClusterAuthentication,
+    authentication: ClusterAuthentication,
 ):
     instance_data = fake_ec2_instance_data(state=ec2_state)
     cluster_instance = create_cluster_from_ec2_instance(