Implemented template validation in the CommandTemplateMap class

streamflow/deployment/template.py

@@ -1,8 +1,21 @@
 from __future__ import annotations
 
-from collections.abc import MutableMapping
+from collections.abc import MutableMapping, MutableSequence
 
-from jinja2 import Template
+from jinja2 import Environment, Template, meta
+
+from streamflow.core.exception import WorkflowDefinitionException
+
+
+def _check_template(name: str, source: str, placeholders: MutableSequence[str]) -> None:
+    ast = Environment(autoescape=True).parse(source)
+    referenced_vars = meta.find_undeclared_variables(ast)
+    for placeholder in placeholders:
+        if placeholder not in referenced_vars:
+            raise WorkflowDefinitionException(
+                f"Template '{name}' does not contain the "
+                f"placeholder '{placeholder}'."
+            )
 
 
 class CommandTemplateMap:
@@ -14,9 +27,11 @@ def __init__(
         self.templates: MutableMapping[str, Template] = {
             "__DEFAULT__": Template(default)
         }
+        _check_template("default", default, ["streamflow_command"])
         if template_map:
             for name, template in template_map.items():
                 self.templates[name] = Template(template)
+                _check_template(name, template, ["streamflow_command"])
 
     def get_command(
         self,

tests/test_connector.py

@@ -13,16 +13,21 @@
 
 from streamflow.core.context import StreamFlowContext
 from streamflow.core.deployment import Connector, ExecutionLocation
-from streamflow.core.exception import WorkflowExecutionException
+from streamflow.core.exception import (
+    WorkflowDefinitionException,
+    WorkflowExecutionException,
+)
 from streamflow.deployment.connector import SSHConnector
 from streamflow.deployment.future import FutureConnector
+from streamflow.deployment.template import CommandTemplateMap
 from tests.conftest import get_class_callables
 from tests.utils.connector import (
     FailureConnector,
     FailureConnectorException,
     SSHChannelErrorConnector,
 )
 from tests.utils.deployment import (
+    get_deployment,
     get_failure_deployment_config,
     get_location,
     get_ssh_deployment_config,
@@ -189,3 +194,81 @@ async def test_ssh_connector_multiple_request_fail(
             )
             is not None
         )
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("deployment_t", ["ssh", "slurm"])
+async def test_templates(
+    chosen_deployment_types: MutableSequence[str],
+    context: StreamFlowContext,
+    deployment_t: str,
+) -> None:
+    """Test command templates"""
+    if deployment_t not in chosen_deployment_types:
+        pytest.skip(f"Deployment {deployment_t} was not activated")
+    deployment = get_deployment(deployment_t)
+    service = "test_template"
+    connector = context.deployment_manager.get_connector(deployment)
+    locations = await connector.get_available_locations(service=service)
+    location = next(iter(locations.values())).location
+    conn = context.deployment_manager.get_connector(location.deployment)
+    stdout, retcode = await conn.run(
+        location, ["true"], capture_output=True, job_name="job_test"
+    )
+    assert retcode == 0
+    assert stdout == f"I am a {deployment_t} service"
+
+
+def test_command_template():
+    """Test command template validation"""
+    t0 = CommandTemplateMap(default="#!/bin/sh\n\n{{streamflow_command}}")
+    assert t0.is_empty()
+    template_map = {
+        "service_a": "module load python\n{{streamflow_command}}",
+        "service_b": "spack load python\ncd {{ streamflow_workdir }}\n{{ streamflow_environment }}\n{{streamflow_command}}",
+        "service_c": "spack load python\n{{ custom_key }}\n{{streamflow_command}}",
+    }
+    t1 = CommandTemplateMap(
+        default="#!/bin/sh\n\n{{streamflow_command}}", template_map=template_map
+    )
+    assert not t1.is_empty()
+    # Test command
+    assert (
+        t1.get_command(
+            "cat /path/to/file", environment={"APP_NAME": "app1"}, workdir="/tmp/abc"
+        )
+        == "#!/bin/sh\n\ncat /path/to/file"
+    )
+    # Test workdir, environment, command
+    assert (
+        t1.get_command(
+            "cat /path/to/file",
+            template="service_b",
+            environment={"APP_NAME": "app1"},
+            workdir="/tmp/abc",
+        )
+        == 'spack load python\ncd /tmp/abc\nexport APP_NAME="app1"\ncat /path/to/file'
+    )
+    # Test custom keys
+    assert (
+        t1.get_command(
+            "cat /path/to/file",
+            template="service_c",
+            environment={"APP_NAME": "app1"},
+            workdir="/tmp/abc",
+            custom_key="hostname",
+        )
+        == "spack load python\nhostname\ncat /path/to/file"
+    )
+
+    # Test validation of default and service templates
+    with pytest.raises(WorkflowDefinitionException):
+        CommandTemplateMap(default="", template_map=template_map)
+    with pytest.raises(WorkflowDefinitionException):
+        CommandTemplateMap(
+            default="#!/bin/sh\n\n{{streamflow_command}}",
+            template_map={
+                "service_a": "module load python\n{{streamflow_command}}",
+                "service_b": "module load python",
+            },
+        )

tests/utils/deployment.py

@@ -30,6 +30,13 @@
 from tests.utils.data import get_data_path
 
 
+def _create_template_file(content: str = "") -> str:
+    content = content or "#!/bin/sh\n\n{{streamflow_command}}"
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f_service:
+        f_service.write(content)
+    return f_service.name
+
+
 def _get_free_tcp_port() -> int:
     """
     Return a free TCP port to be used for publishing services.
@@ -56,7 +63,7 @@ def get_aiotar_deployment_config(tar_format: str) -> DeploymentConfig:
     )
 
 
-def get_deployment(_context: StreamFlowContext, deployment_t: str) -> str:
+def get_deployment(deployment_t: str) -> str:
     match deployment_t:
         case "aiotar":
             return "aiotar"
@@ -215,7 +222,7 @@ def get_local_deployment_config(
 async def get_location(
     _context: StreamFlowContext, deployment_t: str
 ) -> ExecutionLocation:
-    deployment = get_deployment(_context, deployment_t)
+    deployment = get_deployment(deployment_t)
     service = get_service(_context, deployment_t)
     connector = _context.deployment_manager.get_connector(deployment)
     locations = await connector.get_available_locations(service=service)
@@ -286,7 +293,12 @@ async def get_slurm_deployment_config(_context: StreamFlowContext):
         type="slurm",
         config={
             "services": {
-                "test": {"partition": "docker", "nodes": 2, "ntasksPerNode": 1}
+                "test": {"partition": "docker", "nodes": 2, "ntasksPerNode": 1},
+                "test_template": {
+                    "file": _create_template_file(
+                        "#!/bin/sh\n\necho 'I am a slurm service'\n{{streamflow_command}}"
+                    )
+                },
             }
         },
         external=False,
@@ -306,8 +318,8 @@ async def get_ssh_deployment_config(_context: StreamFlowContext):
         key_size=4096,
     )
     public_key = skey.export_public_key().decode("utf-8")
-    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f:
-        skey.write_private_key(f.name)
+    with tempfile.NamedTemporaryFile(mode="w", delete=False) as f_key:
+        skey.write_private_key(f_key.name)
     ssh_port = _get_free_tcp_port()
     docker_config = DeploymentConfig(
         name="linuxserver-ssh-docker",
@@ -327,17 +339,22 @@ async def get_ssh_deployment_config(_context: StreamFlowContext):
         name="linuxserver-ssh",
         type="ssh",
         config={
+            "maxConcurrentSessions": 10,
             "nodes": [
                 {
                     "checkHostKey": False,
                     "hostname": f"127.0.0.1:{ssh_port}",
-                    "sshKey": f.name,
+                    "sshKey": f_key.name,
                     "username": "linuxserver.io",
                     "retries": 2,
                     "retryDelay": 5,
                 }
             ],
-            "maxConcurrentSessions": 10,
+            "services": {
+                "test_template": _create_template_file(
+                    "#!/bin/sh\n\necho 'I am a ssh service'\n{{streamflow_command}}"
+                )
+            },
         },
         workdir="/tmp",
         external=False,