From 530668784a33fcf79c94f9ce4aa9d77e9067f578 Mon Sep 17 00:00:00 2001
From: Anton Mokhovikov <ammokhov@amazon.com>
Date: Wed, 20 Dec 2023 12:11:51 -0800
Subject: [PATCH] Adding regex based exports replacement (#1049)

---
 src/rpdk/core/test.py | 35 +++++++++++++++++++++++++----------
 tests/test_test.py    | 39 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 10 deletions(-)

diff --git a/src/rpdk/core/test.py b/src/rpdk/core/test.py
index 3fefa31a..81d70c30 100644
--- a/src/rpdk/core/test.py
+++ b/src/rpdk/core/test.py
@@ -5,6 +5,7 @@
 import json
 import logging
 import os
+import re
 from argparse import SUPPRESS
 from collections import OrderedDict
 from contextlib import contextmanager
@@ -12,7 +13,6 @@
 from tempfile import NamedTemporaryFile
 
 import pytest
-from jinja2 import Environment, Template, meta
 from jsonschema import Draft6Validator
 from jsonschema.exceptions import ValidationError
 
@@ -116,10 +116,26 @@ def get_cloudformation_exports(region_name, endpoint_url, role_arn, profile_name
     return exports
 
 
-def render_jinja(overrides_string, region_name, endpoint_url, role_arn, profile_name):
-    env = Environment(autoescape=True)
-    parsed_content = env.parse(overrides_string)
-    variables = meta.find_undeclared_variables(parsed_content)
+def _stub_exports(template, exports, pattern):
+    def __retrieve_args(match):
+        try:
+            export = str(match.group(1)).strip()
+            value_to_stub = str(exports[export])
+        except LookupError as e:
+            LOG.error(str(e))
+            raise ValueError(
+                f"Export does not contain provided undeclared variable '{export}'. {e}"
+            ) from e
+        return value_to_stub
+
+    return re.sub(pattern, __retrieve_args, template)
+
+
+def render_template(
+    overrides_string, region_name, endpoint_url, role_arn, profile_name
+):
+    regex = r"{{([-A-Za-z0-9:\s]+?)}}"
+    variables = set(str(match).strip() for match in re.findall(regex, overrides_string))
     if variables:
         exports = get_cloudformation_exports(
             region_name, endpoint_url, role_arn, profile_name
@@ -132,8 +148,7 @@ def render_jinja(overrides_string, region_name, endpoint_url, role_arn, profile_
             )
             LOG.warning(invalid_exports_message, invalid_exports)
             return empty_override()
-        overrides_template = Template(overrides_string)
-        to_return = json.loads(overrides_template.render(exports))
+        to_return = json.loads(_stub_exports(overrides_string, exports, regex))
     else:
         to_return = json.loads(overrides_string)
     return to_return
@@ -158,7 +173,7 @@ def get_overrides(root, region_name, endpoint_url, role_arn, profile_name):
     path = root / "overrides.json"
     try:
         with path.open("r", encoding="utf-8") as f:
-            overrides_raw = render_jinja(
+            overrides_raw = render_template(
                 f.read(), region_name, endpoint_url, role_arn, profile_name
             )
     except FileNotFoundError:
@@ -195,7 +210,7 @@ def get_hook_overrides(root, region_name, endpoint_url, role_arn, profile_name):
     path = root / "overrides.json"
     try:
         with path.open("r", encoding="utf-8") as f:
-            overrides_raw = render_jinja(
+            overrides_raw = render_template(
                 f.read(), region_name, endpoint_url, role_arn, profile_name
             )
     except FileNotFoundError:
@@ -264,7 +279,7 @@ def get_inputs(root, region_name, endpoint_url, value, role_arn, profile_name):
 
                 file_path = path / file
                 with file_path.open("r", encoding="utf-8") as f:
-                    overrides_raw = render_jinja(
+                    overrides_raw = render_template(
                         f.read(), region_name, endpoint_url, role_arn, profile_name
                     )
                 overrides = {}
diff --git a/tests/test_test.py b/tests/test_test.py
index 6f174962..68f13dbe 100644
--- a/tests/test_test.py
+++ b/tests/test_test.py
@@ -22,6 +22,7 @@
     DEFAULT_FUNCTION,
     DEFAULT_PROFILE,
     DEFAULT_REGION,
+    _stub_exports,
     _validate_sam_args,
     empty_hook_override,
     empty_override,
@@ -111,6 +112,44 @@ def _get_expected_marker_options(artifact_type):
     )
 
 
+@pytest.mark.parametrize(
+    "template_string,exports,expected",
+    [
+        (
+            '{"newName": "{{name}}","newLastName": {"newName": "{{ lastname }}"}}',
+            {"name": "Jon", "lastname": "Snow"},
+            '{"newName": "Jon","newLastName": {"newName": "Snow"}}',
+        ),
+        (
+            '{"newName": " {{ name  }}","newLastName": {"newName": "{{ lastname } }"}}',
+            {"name": "Jon", "lastname": "Snow"},
+            '{"newName": " Jon","newLastName": {"newName": "{{ lastname } }"}}',
+        ),
+    ],
+)
+def test_stub_exports(template_string, exports, expected):
+    assert expected == _stub_exports(
+        template_string, exports, r"{{([-A-Za-z0-9:\s]+?)}}"
+    )
+
+
+@pytest.mark.parametrize(
+    "template_string,exports",
+    [
+        (
+            '{"newName": "{{name}}","newLastName": {"newName": "{{ lastname }}"}}',
+            {"name": "Jon"},
+        ),
+    ],
+)
+def test_stub_exports_exception(template_string, exports):
+    with pytest.raises(ValueError) as e:
+        _stub_exports(template_string, exports, r"{{([-A-Za-z0-9:\s]+?)}}")
+        assert (
+            str(e) == "Export does not contain provided undeclared variable 'lastname'"
+        )
+
+
 def create_input_file(base, create_string, update_string, invalid_string):
     path = base / "inputs"
     os.mkdir(path, mode=0o777)