def test_get_variable(self):
param = {'Name': '/airflow/variables/hello', 'Type': 'String', 'Value': 'world'}
ssm_backend = SystemsManagerParameterStoreBackend()
ssm_backend.client.put_parameter(**param)
returned_uri = ssm_backend.get_variable('hello')
assert 'world' == returned_uri
def test_get_variable_secret_string(self):
param = {
'Name': '/airflow/variables/hello',
'Type': 'SecureString',
'Value': 'world'
}
ssm_backend = SystemsManagerParameterStoreBackend()
ssm_backend.client.put_parameter(**param)
returned_uri = ssm_backend.get_variable('hello')
self.assertEqual('world', returned_uri)
def test_get_variable_non_existent_key(self):
"""
Test that if Variable key is not present in SSM,
SystemsManagerParameterStoreBackend.get_variables should return None
"""
param = {'Name': '/airflow/variables/hello', 'Type': 'String', 'Value': 'world'}
ssm_backend = SystemsManagerParameterStoreBackend()
ssm_backend.client.put_parameter(**param)
assert ssm_backend.get_variable("test_mysql") is None
def test_get_config(self):
param = {
'Name': '/airflow/config/sql_alchemy_conn',
'Type': 'String',
'Value': 'sqlite:///Users/test_user/airflow.db',
}
ssm_backend = SystemsManagerParameterStoreBackend()
ssm_backend.client.put_parameter(**param)
returned_uri = ssm_backend.get_config('sql_alchemy_conn')
assert 'sqlite:///Users/test_user/airflow.db' == returned_uri
def test_get_conn_uri(self):
param = {
'Name': '/airflow/connections/test_postgres',
'Type': 'String',
'Value': 'postgresql://*****:*****@host:5432/airflow',
}
ssm_backend = SystemsManagerParameterStoreBackend()
ssm_backend.client.put_parameter(**param)
returned_uri = ssm_backend.get_conn_uri(conn_id="test_postgres")
assert 'postgresql://*****:*****@host:5432/airflow' == returned_uri
def test_config_prefix_none_value(self, mock_get_secret):
"""
Test that if Variable key is not present in SSM,
SystemsManagerParameterStoreBackend.get_config should return None,
SystemsManagerParameterStoreBackend._get_secret should not be called
"""
kwargs = {'config_prefix': None}
ssm_backend = SystemsManagerParameterStoreBackend(**kwargs)
assert ssm_backend.get_config("config") is None
mock_get_secret.assert_not_called()
def test_get_conn_uri_non_existent_key(self):
"""
Test that if the key with connection ID is not present in SSM,
SystemsManagerParameterStoreBackend.get_connections should return None
"""
conn_id = "test_mysql"
param = {
'Name': '/airflow/connections/test_postgres',
'Type': 'String',
'Value': 'postgresql://*****:*****@host:5432/airflow',
}
ssm_backend = SystemsManagerParameterStoreBackend()
ssm_backend.client.put_parameter(**param)
self.assertIsNone(ssm_backend.get_conn_uri(conn_id=conn_id))
self.assertEqual([], ssm_backend.get_connections(conn_id=conn_id))
def test_aws_ssm_get_connections(self, mock_get_uri):
mock_get_uri.return_value = "scheme://*****:*****@host:100"
conn_list = SystemsManagerParameterStoreBackend().get_connections(
"fake_conn")
conn = conn_list[0]
assert conn.host == 'host'
