bhirsz commented on code in PR #24643:
URL: https://github.com/apache/airflow/pull/24643#discussion_r909298070
##########
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:
since fetch_variable defaults default_value to None, we can just do this:
```suggestion
default_value = self.variable_defaults.get(variable, None)
value = fetch_variable(variable, default_value,
test_name=self.test_name)
```
--
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]
