diff --git a/tests/aws/services/stepfunctions/v2/evaluate_jsonata/test_base_evaluate_expressions.py b/tests/aws/services/stepfunctions/v2/evaluate_jsonata/test_base_evaluate_expressions.py
index 4a1a78538a52d..42cc505ba1b19 100644
--- a/tests/aws/services/stepfunctions/v2/evaluate_jsonata/test_base_evaluate_expressions.py
+++ b/tests/aws/services/stepfunctions/v2/evaluate_jsonata/test_base_evaluate_expressions.py
@@ -91,13 +91,7 @@ def test_base_task(
     @pytest.mark.parametrize(
         "expression_dict",
         [
-            pytest.param(
-                {"Items": EJT.JSONATA_ARRAY_ELEMENT_EXPRESSION},
-                marks=pytest.mark.skipif(
-                    condition=not is_aws_cloud(),
-                    reason="Single-quote esacped JSONata expressions are not yet supported",
-                ),
-            ),
+            {"Items": EJT.JSONATA_ARRAY_ELEMENT_EXPRESSION},
             {"Items": EJT.JSONATA_ARRAY_ELEMENT_EXPRESSION_DOUBLE_QUOTES},
             {"MaxConcurrency": EJT.JSONATA_NUMBER_EXPRESSION},
             {"ToleratedFailurePercentage": EJT.JSONATA_NUMBER_EXPRESSION},
diff --git a/tests/unit/services/stepfunctions/test_arguments.py b/tests/unit/services/stepfunctions/test_arguments.py
deleted file mode 100644
index b75c981aadac6..0000000000000
--- a/tests/unit/services/stepfunctions/test_arguments.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# import json
-#
-# import pytest
-#
-# from localstack.services.stepfunctions.templates.querylanguage.query_language_templates import (
-#     QueryLanguageTemplate,
-# )
-# from localstack.services.stepfunctions.asl.component.common.query_language import QueryLanguageMode
-# from localstack.services.stepfunctions.asl.component.program.program import Program
-# from localstack.services.stepfunctions.asl.parse.asl_parser import AmazonStateLanguageParser
-#
-#
-# class TestJSONataIntegration:
-#     # TODO add test cases for MAP and Parallel states, but docs aren't specific enough.
-#
-#     @staticmethod
-#     def _parse_template_file(query_language_template_filepath: str) -> Program:
-#         template = QueryLanguageTemplate.load_sfn_template(query_language_template_filepath)
-#         definition = json.dumps(template)
-#         program: Program = AmazonStateLanguageParser.parse(definition)[0]  # noqa
-#         return program
-#
-#     def test_pass_jsonata(self):
-#         program: Program = self._parse_template_file(QueryLanguageTemplate.BASE_PASS_JSONATA)
-#         assert program.query_language.query_language_mode == QueryLanguageMode.JSONata
-#         assert (
-#             program.states.states["StartState"].query_language.query_language_mode
-#             == QueryLanguageMode.JSONata
-#         )
-#
-#     @pytest.mark.parametrize(
-#         "template_filepath",
-#         [
-#             QueryLanguageTemplate.BASE_PASS_JSONATA_OVERRIDE,
-#             QueryLanguageTemplate.BASE_PASS_JSONATA_OVERRIDE_DEFAULT,
-#         ],
-#         ids=["BASE_PASS_JSONATA_OVERRIDE", "BASE_PASS_JSONATA_OVERRIDE_DEFAULT"],
-#     )
-#     def test_pass_jsonata_override(self, template_filepath):
-#         program: Program = self._parse_template_file(template_filepath)
-#         assert program.query_language.query_language_mode == QueryLanguageMode.JSONPath
-#         assert (
-#             program.states.states["StartState"].query_language.query_language_mode
-#             == QueryLanguageMode.JSONata
-#         )
-#
-#     def test_base_pass_jsonpath(self):
-#         program: Program = self._parse_template_file(QueryLanguageTemplate.BASE_PASS_JSONPATH)
-#         assert program.query_language.query_language_mode == QueryLanguageMode.JSONPath
-#         assert (
-#             program.states.states["StartState"].query_language.query_language_mode
-#             == QueryLanguageMode.JSONPath
-#         )
diff --git a/tests/unit/services/stepfunctions/test_jsonata_integration.py b/tests/unit/services/stepfunctions/test_jsonata_integration.py
deleted file mode 100644
index de47601a063ae..0000000000000
--- a/tests/unit/services/stepfunctions/test_jsonata_integration.py
+++ /dev/null
@@ -1,185 +0,0 @@
-import pytest
-
-from localstack.services.stepfunctions.asl.eval.states import (
-    ContextObjectData,
-    ExecutionData,
-    StateMachineData,
-    States,
-)
-from localstack.services.stepfunctions.asl.jsonata.jsonata import (
-    IllegalJSONataVariableReference,
-    JSONataException,
-    JSONataExpression,
-    VariableDeclarations,
-    VariableReference,
-    compose_jsonata_expression,
-    encode_jsonata_variable_declarations,
-    eval_jsonata_expression,
-    extract_jsonata_variable_references,
-)
-
-POSITIVE_SCENARIOS = [
-    (
-        "base_int_out",
-        """(
-        $x := 10;
-        $x)""",
-        10,
-    ),
-    (
-        "base_float_out",
-        """(
-        $x := 10.1;
-        $x)""",
-        10.1,
-    ),
-    (
-        "base_bool_out",
-        """(
-        $x := true;
-        $x)""",
-        True,
-    ),
-    (
-        "base_array_out",
-        """(
-        $x := [1,2,3];
-        $x)""",
-        [1, 2, 3],
-    ),
-    (
-        "base_object_out",
-        """(
-        $x := {
-            "value": 1
-        };
-        $x)""",
-        {"value": 1},
-    ),
-    (
-        "expression_int_out",
-        """(
-        $obj := {
-            "value": 99
-        };
-        $values := [3, 3];
-        $x := 10;
-        $obj.value+ $sum($values) + $x)""",
-        115,
-    ),
-]
-POSITIVE_SCENARIOS_IDS = [scenario[0] for scenario in POSITIVE_SCENARIOS]
-POSITIVE_SCENARIOS_EXPR_OUTPUT_PARIS = [tuple(scenario[1:]) for scenario in POSITIVE_SCENARIOS]
-
-
-NEGATIVE_SCENARIOS = [("null-input", None), ("empty-input", ""), ("syntax-error-semi", ";")]
-NEGATIVE_SCENARIOS_IDS = [scenario[0] for scenario in NEGATIVE_SCENARIOS]
-NEGATIVE_SCENARIOS_EXPRESSIONS = [scenario[1] for scenario in NEGATIVE_SCENARIOS]
-
-
-VARIABLE_ASSIGNMENT_ENCODING = [
-    ("int", {"$var1": 3}, "$var1:=3;"),
-    ("float", {"$var1": 3.2}, "$var1:=3.2;"),
-    ("null", {"$var1": None}, "$var1:=null;"),
-    ("string", {"$var1": "string_lit"}, '$var1:="string_lit";'),
-    ("list", {"$var1": [3, 3.2, None, "string_lit", []]}, '$var1:=[3,3.2,null,"string_lit",[]];'),
-    (
-        "obj",
-        {"$var1": {"string_lit": "string_lit_value"}},
-        '$var1:={"string_lit":"string_lit_value"};',
-    ),
-    (
-        "mult",
-        {
-            "$var0": 0,
-            "$var1": {"string_lit": "string_lit_value"},
-            "$var2": [3, 3.2, None, "string_lit", []],
-        },
-        '$var0:=0;$var1:={"string_lit":"string_lit_value"};$var2:=[3,3.2,null,"string_lit",[]];',
-    ),
-]
-VARIABLE_ASSIGNMENT_ENCODING_IDS = [scenario[0] for scenario in VARIABLE_ASSIGNMENT_ENCODING]
-VARIABLE_ASSIGNMENT_ENCODING_SCENARIOS = [
-    tuple(scenario[1:]) for scenario in VARIABLE_ASSIGNMENT_ENCODING
-]
-
-STATES_ACCESSES_STATES = States(
-    context=ContextObjectData(
-        Execution=ExecutionData(
-            Id="test-exec-arn",
-            Input={"items": [None, 1, 1.1, True, [], {"key": "string_lit"}]},
-            Name="test-name",
-            RoleArn="test-role",
-            StartTime="test-start-time",
-        ),
-        StateMachine=StateMachineData(Id="test-arn", Name="test-name"),
-    )
-)
-STATES_ACCESSES_STATES.set_result({"result_key": "result_value"})
-STATES_ACCESSES = [
-    ("input", "$states.input", {"items": [None, 1, 1.1, True, [], {"key": "string_lit"}]}),
-    ("input.items", "$states.input.items", [None, 1, 1.1, True, [], {"key": "string_lit"}]),
-    (
-        "input.items-result",
-        "[$states.input.items, $states.result]",
-        [None, 1, 1.1, True, [], {"key": "string_lit"}, {"result_key": "result_value"}],
-    ),
-]
-STATES_ACCESSES_IDS = [scenario[0] for scenario in STATES_ACCESSES]
-STATES_ACCESSES_SCENARIOS = [tuple(scenario[1:]) for scenario in STATES_ACCESSES]
-
-
-class TestJSONataIntegration:
-    @pytest.mark.parametrize(
-        "expression, expected",
-        POSITIVE_SCENARIOS_EXPR_OUTPUT_PARIS,
-        ids=POSITIVE_SCENARIOS_IDS,
-    )
-    def test_expressions_positive(self, expression, expected):
-        result = eval_jsonata_expression(expression)
-        assert result == expected
-
-    @pytest.mark.parametrize(
-        "expression",
-        NEGATIVE_SCENARIOS_EXPRESSIONS,
-        ids=NEGATIVE_SCENARIOS_IDS,
-    )
-    def test_expressions_negative(self, expression):
-        with pytest.raises(JSONataException):
-            eval_jsonata_expression(expression)
-
-    def test_variable_assignment_extraction_positive(self):
-        expression = "$a;$a0;$a0_;$a_0;$_a;$var1.var2.var3;$var$;$va$r;$_0a$.b$0;$var$$;$va$r$$.b$$"
-        references = extract_jsonata_variable_references(expression)
-        assert sorted(references) == sorted(expression.split(";"))
-
-    def test_variable_assignment_extraction_negative(self):
-        illegal_expressions = ["$", "$$"]
-        for illegal_expression in illegal_expressions:
-            with pytest.raises(IllegalJSONataVariableReference):
-                extract_jsonata_variable_references(illegal_expression)
-
-    @pytest.mark.parametrize(
-        "bindings, expected",
-        VARIABLE_ASSIGNMENT_ENCODING_SCENARIOS,
-        ids=VARIABLE_ASSIGNMENT_ENCODING_IDS,
-    )
-    def test_variable_assignment_encoding(self, bindings, expected):
-        encoding = encode_jsonata_variable_declarations(bindings)
-        assert encoding == expected
-
-    @pytest.mark.parametrize(
-        "expression, expected", STATES_ACCESSES_SCENARIOS, ids=STATES_ACCESSES_IDS
-    )
-    def test_states_access(self, expression, expected):
-        variable_references: set[VariableReference] = extract_jsonata_variable_references(
-            expression
-        )
-        variable_declarations: VariableDeclarations = (
-            STATES_ACCESSES_STATES.to_variable_declarations(variable_references=variable_references)
-        )
-        rich_jsonata_expression: JSONataExpression = compose_jsonata_expression(
-            final_jsonata_expression=expression, variable_declarations_list=[variable_declarations]
-        )
-        result = eval_jsonata_expression(rich_jsonata_expression)
-        assert result == expected
diff --git a/tests/unit/services/stepfunctions/test_jsonata_payload.py b/tests/unit/services/stepfunctions/test_jsonata_payload.py
deleted file mode 100644
index f1ed278830a01..0000000000000
--- a/tests/unit/services/stepfunctions/test_jsonata_payload.py
+++ /dev/null
@@ -1,122 +0,0 @@
-import pytest
-from antlr4 import CommonTokenStream, InputStream
-
-from localstack.aws.api.stepfunctions import StateMachineType
-from localstack.services.stepfunctions.asl.antlr.runtime.ASLLexer import ASLLexer
-from localstack.services.stepfunctions.asl.antlr.runtime.ASLParser import ASLParser
-from localstack.services.stepfunctions.asl.component.common.assign.assign_template_value_object import (
-    AssignTemplateValueObject,
-)
-from localstack.services.stepfunctions.asl.component.common.query_language import (
-    QueryLanguage,
-    QueryLanguageMode,
-)
-from localstack.services.stepfunctions.asl.eval.environment import Environment
-from localstack.services.stepfunctions.asl.eval.evaluation_details import AWSExecutionDetails
-from localstack.services.stepfunctions.asl.eval.event.event_manager import EventHistoryContext
-from localstack.services.stepfunctions.asl.eval.program_state import ProgramRunning
-from localstack.services.stepfunctions.asl.eval.states import (
-    ContextObjectData,
-    ExecutionData,
-    StateMachineData,
-)
-from localstack.services.stepfunctions.asl.parse.preprocessor import Preprocessor
-
-POSITIVE_SCENARIOS = [
-    (
-        "base_string_bindings",
-        """{
-            "jsonataexpr": "{% ($name := 'NameString'; $name) %}",
-            "stringlit1": " {% ($name := 'NameString'; $name) %}",
-            "stringlit2": "{% ($name := 'NameString'; $name) %} ",
-            "stringlit3": "stringlit",
-            "stringlig4": "$.stringlit",
-            "stringlig5": "$.stringlit"
-        }""",
-        {
-            "jsonataexpr": "NameString",
-            "stringlit1": " {% ($name := 'NameString'; $name) %}",
-            "stringlit2": "{% ($name := 'NameString'; $name) %} ",
-            "stringlit3": "stringlit",
-            "stringlig4": "$.stringlit",
-            "stringlig5": "$.stringlit",
-        },
-    ),
-    (
-        "base_types",
-        """{
-            "jsonataexpr": "{% ($name := 'namestring'; $name) %}",
-            "null": null,
-            "int": 1,
-            "float": 0.1,
-            "boolt": true,
-            "boolf": false,
-            "arr": [null, 1, 0.1, true, [], {"jsonataexpr": "{% ($name := 'namestring'; $name) %}"}],
-            "obj": {"jsonataexpr": "{% ($name := 'namestring'; $name) %}"}
-        }""",
-        {
-            "jsonataexpr": "namestring",
-            "null": None,
-            "int": 1,
-            "float": 0.1,
-            "boolf": False,
-            "boolt": True,
-            "arr": [None, 1, 0.1, True, [], {"jsonataexpr": "namestring"}],
-            "obj": {"jsonataexpr": "namestring"},
-        },
-    ),
-]
-POSITIVE_SCENARIOS_IDS = [scenario[0] for scenario in POSITIVE_SCENARIOS]
-POSITIVE_SCENARIOS_EXPR_OUTPUT_PARIS = [tuple(scenario[1:]) for scenario in POSITIVE_SCENARIOS]
-
-
-def parse_payload(payload_derivation: str) -> AssignTemplateValueObject:
-    input_stream = InputStream(payload_derivation)
-    lexer = ASLLexer(input_stream)
-    stream = CommonTokenStream(lexer)
-    parser = ASLParser(stream)
-    tree = parser.assign_template_value_object()
-    preprocessor = Preprocessor()
-    # simulate a jsonata query language top level definition.
-    preprocessor._query_language_per_scope.append(
-        QueryLanguage(query_language_mode=QueryLanguageMode.JSONata)
-    )
-    jsonata_payload_object = preprocessor.visit(tree)
-    preprocessor._query_language_per_scope.clear()
-    return jsonata_payload_object
-
-
-def evaluate_payload(jsonata_payload_object: AssignTemplateValueObject) -> dict:
-    env = Environment(
-        aws_execution_details=AWSExecutionDetails("test-account", "test-region", "test-role"),
-        execution_type=StateMachineType.STANDARD,
-        context=ContextObjectData(
-            Execution=ExecutionData(
-                Id="test-exec-arn",
-                Input=dict(),
-                Name="test-name",
-                RoleArn="test-role",
-                StartTime="test-start-time",
-            ),
-            StateMachine=StateMachineData(Id="test-arn", Name="test-name"),
-        ),
-        event_history_context=EventHistoryContext.of_program_start(),
-        cloud_watch_logging_session=None,
-        activity_store=dict(),
-    )
-    env._program_state = ProgramRunning()
-    env.next_state_name = "test-state"
-    jsonata_payload_object.eval(env)
-    return env.stack.pop()
-
-
-class TestJSONataPayload:
-    @pytest.mark.parametrize(
-        "derivation, output",
-        POSITIVE_SCENARIOS_EXPR_OUTPUT_PARIS,
-        ids=POSITIVE_SCENARIOS_IDS,
-    )
-    def test_derivation_positive(self, derivation, output):
-        jsonata_payload_object = parse_payload(derivation)
-        result = evaluate_payload(jsonata_payload_object)
-        assert result == output
diff --git a/tests/unit/services/stepfunctions/test_query_language_parsing.py b/tests/unit/services/stepfunctions/test_query_language_parsing.py
deleted file mode 100644
index d67f8a6161515..0000000000000
--- a/tests/unit/services/stepfunctions/test_query_language_parsing.py
+++ /dev/null
@@ -1,53 +0,0 @@
-# import json
-#
-# import pytest
-#
-# from aws.services.stepfunctions.templates.querylanguage.query_language_templates import (
-#     QueryLanguageTemplate,
-# )
-# from localstack.services.stepfunctions.asl.component.common.query_language import QueryLanguageMode
-# from localstack.services.stepfunctions.asl.component.program.program import Program
-# from localstack.services.stepfunctions.asl.parse.asl_parser import AmazonStateLanguageParser
-#
-#
-# class TestJSONataIntegration:
-#     # TODO add test cases for MAP and Parallel states, but docs aren't specific enough.
-#
-#     @staticmethod
-#     def _parse_template_file(query_language_template_filepath: str) -> Program:
-#         template = QueryLanguageTemplate.load_sfn_template(query_language_template_filepath)
-#         definition = json.dumps(template)
-#         program: Program = AmazonStateLanguageParser.parse(definition)[0]  # noqa
-#         return program
-#
-#     def test_pass_jsonata(self):
-#         program: Program = self._parse_template_file(QueryLanguageTemplate.BASE_PASS_JSONATA)
-#         assert program.query_language.query_language_mode == QueryLanguageMode.JSONata
-#         assert (
-#             program.states.states["StartState"].query_language.query_language_mode
-#             == QueryLanguageMode.JSONata
-#         )
-#
-#     @pytest.mark.parametrize(
-#         "template_filepath",
-#         [
-#             QueryLanguageTemplate.BASE_PASS_JSONATA_OVERRIDE,
-#             QueryLanguageTemplate.BASE_PASS_JSONATA_OVERRIDE_DEFAULT,
-#         ],
-#         ids=["BASE_PASS_JSONATA_OVERRIDE", "BASE_PASS_JSONATA_OVERRIDE_DEFAULT"],
-#     )
-#     def test_pass_jsonata_override(self, template_filepath):
-#         program: Program = self._parse_template_file(template_filepath)
-#         assert program.query_language.query_language_mode == QueryLanguageMode.JSONPath
-#         assert (
-#             program.states.states["StartState"].query_language.query_language_mode
-#             == QueryLanguageMode.JSONata
-#         )
-#
-#     def test_base_pass_jsonpath(self):
-#         program: Program = self._parse_template_file(QueryLanguageTemplate.BASE_PASS_JSONPATH)
-#         assert program.query_language.query_language_mode == QueryLanguageMode.JSONPath
-#         assert (
-#             program.states.states["StartState"].query_language.query_language_mode
-#             == QueryLanguageMode.JSONPath
-#         )
diff --git a/tests/unit/services/stepfunctions/test_usage_metrics_static_analyser.py b/tests/unit/services/stepfunctions/test_usage_metrics_static_analyser.py
index fc1e457d477be..86779bebb1096 100644
--- a/tests/unit/services/stepfunctions/test_usage_metrics_static_analyser.py
+++ b/tests/unit/services/stepfunctions/test_usage_metrics_static_analyser.py
@@ -5,33 +5,96 @@
 from localstack.services.stepfunctions.asl.static_analyser.usage_metrics_static_analyser import (
     UsageMetricsStaticAnalyser,
 )
-from tests.aws.services.stepfunctions.templates.assign.assign_templates import (
-    AssignTemplate,
+
+BASE_PASS_JSONATA = json.dumps(
+    {
+        "QueryLanguage": "JSONata",
+        "StartAt": "StartState",
+        "States": {
+            "StartState": {"Type": "Pass", "End": True},
+        },
+    }
 )
-from tests.aws.services.stepfunctions.templates.querylanguage.query_language_templates import (
-    QueryLanguageTemplate,
+
+BASE_PASS_JSONPATH = json.dumps(
+    {
+        "QueryLanguage": "JSONPath",
+        "StartAt": "StartState",
+        "States": {
+            "StartState": {"Type": "Pass", "End": True},
+        },
+    }
 )
 
+BASE_PASS_JSONATA_OVERRIDE = json.dumps(
+    {
+        "QueryLanguage": "JSONPath",
+        "StartAt": "StartState",
+        "States": {
+            "StartState": {"QueryLanguage": "JSONata", "Type": "Pass", "End": True},
+        },
+    }
+)
 
-class TestUsageMetricsStaticAnalyser:
-    @staticmethod
-    def _get_query_language_definition(query_language_template_filepath: str) -> str:
-        template = QueryLanguageTemplate.load_sfn_template(query_language_template_filepath)
-        definition = json.dumps(template)
-        return definition
+BASE_PASS_JSONATA_OVERRIDE_DEFAULT = json.dumps(
+    {
+        "StartAt": "StartState",
+        "States": {
+            "StartState": {"QueryLanguage": "JSONata", "Type": "Pass", "End": True},
+        },
+    }
+)
+
+JSONPATH_TO_JSONATA_DATAFLOW = json.dumps(
+    {
+        "StartAt": "StateJsonPath",
+        "States": {
+            "StateJsonPath": {"Type": "Pass", "Assign": {"var": 42}, "Next": "StateJsonata"},
+            "StateJsonata": {
+                "QueryLanguage": "JSONata",
+                "Type": "Pass",
+                "Output": "{% $var %}",
+                "End": True,
+            },
+        },
+    }
+)
+
+ASSIGN_BASE_EMPTY = json.dumps(
+    {"StartAt": "State0", "States": {"State0": {"Type": "Pass", "Assign": {}, "End": True}}}
+)
+
+ASSIGN_BASE_SCOPE_MAP = json.dumps(
+    {
+        "StartAt": "State0",
+        "States": {
+            "State0": {
+                "Type": "Map",
+                "ItemProcessor": {
+                    "ProcessorConfig": {"Mode": "INLINE"},
+                    "StartAt": "Inner",
+                    "States": {
+                        "Inner": {
+                            "Type": "Pass",
+                            "Assign": {},
+                            "End": True,
+                        },
+                    },
+                },
+                "End": True,
+            }
+        },
+    }
+)
 
-    @staticmethod
-    def _get_variable_sampling_definition(variable_sampling_template_filepath: str) -> str:
-        template = AssignTemplate.load_sfn_template(variable_sampling_template_filepath)
-        definition = json.dumps(template)
-        return definition
 
+class TestUsageMetricsStaticAnalyser:
     @pytest.mark.parametrize(
-        "template_filepath",
+        "definition",
         [
-            QueryLanguageTemplate.BASE_PASS_JSONATA,
-            QueryLanguageTemplate.BASE_PASS_JSONATA_OVERRIDE,
-            QueryLanguageTemplate.BASE_PASS_JSONATA_OVERRIDE_DEFAULT,
+            BASE_PASS_JSONATA,
+            BASE_PASS_JSONATA_OVERRIDE,
+            BASE_PASS_JSONATA_OVERRIDE_DEFAULT,
         ],
         ids=[
             "BASE_PASS_JSONATA",
@@ -39,58 +102,37 @@ def _get_variable_sampling_definition(variable_sampling_template_filepath: str)
             "BASE_PASS_JSONATA_OVERRIDE_DEFAULT",
         ],
     )
-    def test_jsonata(self, template_filepath):
-        definition = self._get_query_language_definition(template_filepath)
+    def test_jsonata(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
         assert analyser.has_jsonata
         assert not analyser.has_variable_sampling
 
-    @pytest.mark.parametrize(
-        "template_filepath",
-        [
-            QueryLanguageTemplate.BASE_PASS_JSONPATH,
-        ],
-        ids=["BASE_PASS_JSONPATH"],
-    )
-    def test_jsonpath(self, template_filepath):
-        definition = self._get_query_language_definition(template_filepath)
+    @pytest.mark.parametrize("definition", [BASE_PASS_JSONPATH], ids=["BASE_PASS_JSONPATH"])
+    def test_jsonpath(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
         assert not analyser.has_jsonata
         assert not analyser.has_variable_sampling
 
     @pytest.mark.parametrize(
-        "template_filepath",
-        [
-            QueryLanguageTemplate.JSONPATH_TO_JSONATA_DATAFLOW,
-            QueryLanguageTemplate.JSONPATH_ASSIGN_JSONATA_REF,
-        ],
-        ids=["JSONPATH_TO_JSONATA_DATAFLOW", "JSONPATH_ASSIGN_JSONATA_REF"],
+        "definition", [JSONPATH_TO_JSONATA_DATAFLOW], ids=["JSONPATH_TO_JSONATA_DATAFLOW"]
     )
-    def test_jsonata_and_variable_sampling(self, template_filepath):
-        definition = self._get_query_language_definition(template_filepath)
+    def test_jsonata_and_variable_sampling(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
         assert analyser.has_jsonata
         assert analyser.has_variable_sampling
 
     @pytest.mark.parametrize(
-        "template_filepath",
+        "definition",
         [
-            AssignTemplate.BASE_EMPTY,
-            AssignTemplate.BASE_PATHS,
-            AssignTemplate.BASE_SCOPE_MAP,
-            AssignTemplate.BASE_ASSIGN_FROM_LAMBDA_TASK_RESULT,
-            AssignTemplate.BASE_REFERENCE_IN_LAMBDA_TASK_FIELDS,
+            ASSIGN_BASE_EMPTY,
+            ASSIGN_BASE_SCOPE_MAP,
         ],
         ids=[
-            "BASE_EMPTY",
-            "BASE_PATHS",
-            "BASE_SCOPE_MAP",
-            "BASE_ASSIGN_FROM_LAMBDA_TASK_RESULT",
-            "BASE_REFERENCE_IN_LAMBDA_TASK_FIELDS",
+            "ASSIGN_BASE_EMPTY",
+            "ASSIGN_BASE_SCOPE_MAP",
         ],
     )
-    def test_jsonpath_and_variable_sampling(self, template_filepath):
-        definition = self._get_query_language_definition(template_filepath)
+    def test_jsonpath_and_variable_sampling(self, definition):
         analyser = UsageMetricsStaticAnalyser.process(definition)
         assert not analyser.has_jsonata
         assert analyser.has_variable_sampling