[GitHub] [airflow] o-nikolas commented on a diff in pull request #24643: Convert Step Functions Example DAG to System Test (AIP-47)

[GitBox <gi...@apache.org>]

o-nikolas commented on code in PR #24643:
URL: https://github.com/apache/airflow/pull/24643#discussion_r913321709


##########
tests/system/providers/amazon/aws/utils/__init__.py:
##########
@@ -102,7 +105,54 @@ def _fetch_from_ssm(key: str) -> str:
     return value
 
 
-def fetch_variable(key: str, default_value: Optional[str] = None) -> str:
+class SystemTestContextBuilder:
+    """This builder class ultimately constructs a TaskFlow task which is run at
+    runtime (task execution time). This task generates and stores the test ENV_ID as well
+    as any external resources requested (e.g.g IAM Roles, VPC, etc)"""
+
+    def __init__(self):
+        self.variables = []
+        self.variable_defaults = {}
+        self.test_name = _get_test_name()
+        self.env_id = set_env_id()
+
+    def add_variable(self, variable_name: str, **kwargs):
+        """Register a variable to fetch from environment or cloud parameter store"""
+        self.variables.append(variable_name)
+        # default_value is accepted via kwargs so that it is completely optional and no
+        # default value needs to be provided in the method stub (otherwise we wouldn't
+        # be able to tell the difference between our default value and one provided by
+        # the caller)
+        if 'default_value' in kwargs:
+            self.variable_defaults[variable_name] = kwargs['default_value']
+
+        return self  # Builder recipe; returning self allows chaining
+
+    def build(self):
+        """Build and return a TaskFlow task which will create an env_id and
+        fetch requested variables. Storing everything in xcom for downstream
+        tasks to use."""
+
+        @task
+        def variable_fetcher(**kwargs):
+            ti = kwargs['ti']
+            for variable in self.variables:
+                if variable in self.variable_defaults:
+                    value = fetch_variable(
+                        variable, self.variable_defaults[variable], test_name=self.test_name
+                    )
+                else:
+                    value = fetch_variable(variable, test_name=self.test_name)

Review Comment:
   This code was written in such a way that any default could be provided, even `None` (since the user may want literally `None` to be the default value), but you're right that the `fetch_variable` code wasn't written with that philosophy so we can go with your suggestion for now until someone requires otherwise.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org