From d7d6c82e83c76e00fa3575d819a9fa320be58424 Mon Sep 17 00:00:00 2001 From: sanderegg <35365065+sanderegg@users.noreply.github.com> Date: Mon, 2 Dec 2024 17:34:54 +0100 Subject: [PATCH] fixed tests in ck --- .../modules/dask.py | 10 +++------- .../utils/clusters.py | 4 ++-- .../utils/dask.py | 4 ++-- .../tests/unit/test_modules_dask.py | 16 +++++++++++----- .../tests/unit/test_utils_clusters.py | 8 ++++---- 5 files changed, 22 insertions(+), 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 index af1d0df0e663..0641e812777b 100644 --- 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 index a6ecfdb81894..5a9402ba0930 100644 --- 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 index 266557358b7f..6dc6a452fe42 100644 --- 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 index 7f0408d7057c..97c831ea789b 100644 --- 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 index 55190cb46a18..96983dd34d5e 100644 --- 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,7 +55,7 @@ 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 @@ -63,7 +63,7 @@ def backend_cluster_auth( 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(