From aba9dedd491fc5949ffb78f3864880bd19120f94 Mon Sep 17 00:00:00 2001
From: Marcin Romaszewicz <marcinromaszewicz@deepmap.ai>
Date: Thu, 24 Mar 2022 12:08:17 -0700
Subject: [PATCH] Allow variable lookups in config dict keys

Sometimes, dict keys need to be generated via lookup, such as for
federated IAM policies, where the conditions need to specify
outputs of other stacks.

Stacker supported lookup substitution inside dictionary values, and
this change extends that behavior to dictionary keys.
---
 stacker/config/__init__.py    | 3 ++-
 stacker/lookups/__init__.py   | 4 +++-
 stacker/tests/test_config.py  | 3 +++
 stacker/tests/test_lookups.py | 7 +++++++
 4 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/stacker/config/__init__.py b/stacker/config/__init__.py
index d14ec49ce..09494b6d5 100644
--- a/stacker/config/__init__.py
+++ b/stacker/config/__init__.py
@@ -166,7 +166,8 @@ def isstr(s):
     elif isinstance(root, dict):
         result = {}
         for k, v in root.items():
-            result[k] = substitute_references(v, environment, exp, full_exp)
+            new_k = substitute_references(k, environment, exp, full_exp)
+            result[new_k] = substitute_references(v, environment, exp, full_exp)
         return result
     elif isstr(root):
         # Strings are the special type where all substitutions happen. If we
diff --git a/stacker/lookups/__init__.py b/stacker/lookups/__init__.py
index e8513a560..4f81fc2d8 100644
--- a/stacker/lookups/__init__.py
+++ b/stacker/lookups/__init__.py
@@ -55,12 +55,14 @@ def extract_lookups(value):
 
     """
     lookups = set()
+
     if isinstance(value, basestring):
         lookups = lookups.union(extract_lookups_from_string(value))
     elif isinstance(value, list):
         for v in value:
             lookups = lookups.union(extract_lookups(v))
     elif isinstance(value, dict):
-        for v in value.values():
+        for k, v in value.items():
+            lookups = lookups.union(extract_lookups(k))
             lookups = lookups.union(extract_lookups(v))
     return lookups
diff --git a/stacker/tests/test_config.py b/stacker/tests/test_config.py
index 44b172be5..9bb968f4d 100644
--- a/stacker/tests/test_config.py
+++ b/stacker/tests/test_config.py
@@ -67,6 +67,8 @@ def test_render_yaml(self):
             substr: prefix-${str_1}-suffix
             multiple: ${str_1}-${str_2}
             dont_match_this: ${output something}
+            key_dict:
+                ${namespace}: foo
         """
         env = """
             namespace: test
@@ -108,6 +110,7 @@ def test_render_yaml(self):
         self.assertEquals(pc['substr'], 'prefix-another str-suffix')
         self.assertEquals(pc['multiple'], 'another str-hello')
         self.assertEquals(pc['dont_match_this'], '${output something}')
+        self.assertEquals(pc['key_dict']["test"], 'foo')
 
     def test_render_yaml_errors(self):
         # We shouldn't be able to substitute an object into a string
diff --git a/stacker/tests/test_lookups.py b/stacker/tests/test_lookups.py
index 444f776c7..3f6341e9d 100644
--- a/stacker/tests/test_lookups.py
+++ b/stacker/tests/test_lookups.py
@@ -35,6 +35,13 @@ def test_lookups_dict(self):
         })
         self.assertEqual(len(lookups), 1)
 
+        # Lookups should work in keys as well.
+        lookups = extract_lookups({
+            "${output fakeStack::FakeKeyName}-something": "${output fakeStack::FakeOutput}",
+            "other": "value",
+        })
+        self.assertEqual(len(lookups), 2)
+
     def test_lookups_mixed(self):
         lookups = extract_lookups({
             "something": "${output fakeStack::FakeOutput}",