added test to be sure the custom boot script comes in

ITISFoundation · Dec 6, 2023 · b8ccae6 · b8ccae6
1 parent 4fd3618
commit b8ccae6
Show file tree

Hide file tree

Showing 4 changed files with 33 additions and 8 deletions.
diff --git a/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters.py b/services/clusters-keeper/src/simcore_service_clusters_keeper/modules/clusters.py
@@ -72,6 +72,7 @@ async def create_cluster(
             get_cluster_name(
                 app_settings, user_id=user_id, wallet_id=wallet_id, is_manager=False
             ),
+            ec2_instance_boot_specs,
         ),
         ami_id=ec2_instance_boot_specs.ami_id,
         key_name=app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_KEY_NAME,

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
@@ -4,7 +4,7 @@
 import json
 from typing import Any, Final
 
-from aws_library.ec2.models import EC2InstanceData
+from aws_library.ec2.models import EC2InstanceBootSpecific, EC2InstanceData
 from fastapi.encoders import jsonable_encoder
 from models_library.api_schemas_clusters_keeper.clusters import (
     ClusterState,
@@ -31,7 +31,9 @@ def _docker_compose_yml_base64_encoded() -> str:
 
 
 def create_startup_script(
-    app_settings: ApplicationSettings, cluster_machines_name_prefix: str
+    app_settings: ApplicationSettings,
+    cluster_machines_name_prefix: str,
+    ec2_boot_specific: EC2InstanceBootSpecific,
 ) -> str:
     assert app_settings.CLUSTERS_KEEPER_EC2_ACCESS  # nosec
     assert app_settings.CLUSTERS_KEEPER_WORKERS_EC2_INSTANCES  # nosec
@@ -59,13 +61,15 @@ def _convert_to_env_dict(entries: dict[str, Any]) -> str:
         f"LOG_LEVEL={app_settings.LOG_LEVEL}",
     ]
 
-    return "\n".join(
+    startup_commands = ec2_boot_specific.custom_boot_scripts.copy()
+    startup_commands.extend(
         [
             f"echo '{_docker_compose_yml_base64_encoded()}' | base64 -d > docker-compose.yml",
             "docker swarm init",
             f"{' '.join(environment_variables)} docker stack deploy --with-registry-auth --compose-file=docker-compose.yml dask_stack",
         ]
     )
+    return "\n".join(startup_commands)
 
 
 def _convert_ec2_state_to_cluster_state(

diff --git a/services/clusters-keeper/tests/unit/conftest.py b/services/clusters-keeper/tests/unit/conftest.py
@@ -106,9 +106,11 @@ def app_environment(
             "PRIMARY_EC2_INSTANCES_SUBNET_ID": faker.pystr(),
             "PRIMARY_EC2_INSTANCES_ALLOWED_TYPES": json.dumps(
                 {
-                    random.choice(ec2_instances): random.choice(  # noqa: S311
-                        EC2InstanceBootSpecific.Config.schema_extra["examples"]
-                    )
+                    random.choice(  # noqa: S311
+                        ec2_instances
+                    ): EC2InstanceBootSpecific.Config.schema_extra["examples"][
+                        1
+                    ]  # NOTE: we use example with custom script
                 }
             ),
             "CLUSTERS_KEEPER_WORKERS_EC2_INSTANCES": "{}",

diff --git a/services/clusters-keeper/tests/unit/test_utils_clusters.py b/services/clusters-keeper/tests/unit/test_utils_clusters.py
@@ -7,7 +7,7 @@
 from typing import Any
 
 import pytest
-from aws_library.ec2.models import EC2InstanceData
+from aws_library.ec2.models import EC2InstanceBootSpecific, EC2InstanceData
 from faker import Faker
 from models_library.api_schemas_clusters_keeper.clusters import ClusterState
 from pytest_simcore.helpers.utils_envs import EnvVarsDict
@@ -24,16 +24,34 @@ def cluster_machines_name_prefix(faker: Faker) -> str:
     return faker.pystr()
 
 
+@pytest.fixture
+def ec2_boot_specs(app_settings: ApplicationSettings) -> EC2InstanceBootSpecific:
+    assert app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES
+    ec2_boot_specs = next(
+        iter(
+            app_settings.CLUSTERS_KEEPER_PRIMARY_EC2_INSTANCES.PRIMARY_EC2_INSTANCES_ALLOWED_TYPES.values()
+        )
+    )
+    assert isinstance(ec2_boot_specs, EC2InstanceBootSpecific)
+    return ec2_boot_specs
+
+
 def test_create_startup_script(
     disabled_rabbitmq: None,
     mocked_ec2_server_envs: EnvVarsDict,
     mocked_redis_server: None,
     app_settings: ApplicationSettings,
     cluster_machines_name_prefix: str,
     clusters_keeper_docker_compose: dict[str, Any],
+    ec2_boot_specs: EC2InstanceBootSpecific,
 ):
-    startup_script = create_startup_script(app_settings, cluster_machines_name_prefix)
+    startup_script = create_startup_script(
+        app_settings, cluster_machines_name_prefix, ec2_boot_specs
+    )
     assert isinstance(startup_script, str)
+    assert len(ec2_boot_specs.custom_boot_scripts) > 0
+    for boot_script in ec2_boot_specs.custom_boot_scripts:
+        assert boot_script in startup_script
     # we have commands to pipe into a docker-compose file
     assert " | base64 -d > docker-compose.yml" in startup_script
     # we have commands to init a docker-swarm