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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
