def test_get_variable(self): param = {'SecretId': 'airflow/variables/hello', 'SecretString': 'world'} secrets_manager_backend = SecretsManagerBackend() secrets_manager_backend.client.put_secret_value(**param) returned_uri = secrets_manager_backend.get_variable('hello') assert 'world' == returned_uri
def test_get_variable_non_existent_key(self): """ Test that if Variable key is not present, SystemsManagerParameterStoreBackend.get_variables should return None """ param = {'SecretId': 'airflow/variables/hello', 'SecretString': 'world'} secrets_manager_backend = SecretsManagerBackend() secrets_manager_backend.client.put_secret_value(**param) assert secrets_manager_backend.get_variable("test_mysql") is None
def test_get_conn_uri(self): param = { 'SecretId': 'airflow/connections/test_postgres', 'SecretString': 'postgresql://*****:*****@host:5432/airflow', } secrets_manager_backend = SecretsManagerBackend() secrets_manager_backend.client.put_secret_value(**param) returned_uri = secrets_manager_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 AWS Secrets Manager, SecretsManagerBackend.get_config should return None, SecretsManagerBackend._get_secret should not be called """ kwargs = {'config_prefix': None} secrets_manager_backend = SecretsManagerBackend(**kwargs) self.assertIsNone(secrets_manager_backend.get_config("config")) mock_get_secret.assert_not_called()
def test_connection_prefix_none_value(self, mock_get_secret): """ Test that if Variable key is not present in AWS Secrets Manager, SecretsManagerBackend.get_conn_uri should return None, SecretsManagerBackend._get_secret should not be called """ kwargs = {'connections_prefix': None} secrets_manager_backend = SecretsManagerBackend(**kwargs) assert secrets_manager_backend.get_conn_uri("test_mysql") is None mock_get_secret.assert_not_called()
def test_variable_prefix_none_value(self, mock_get_secret): """ Test that if Variable key is not present in AWS Secrets Manager, SecretsManagerBackend.get_variables should return None, SecretsManagerBackend._get_secret should not be called """ kwargs = {'variables_prefix': None} secrets_manager_backend = SecretsManagerBackend(**kwargs) assert secrets_manager_backend.get_variable("hello") 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, SecretsManagerBackend.get_connections should return None """ conn_id = "test_mysql" param = { 'SecretId': 'airflow/connections/test_postgres', 'SecretString': 'postgresql://*****:*****@host:5432/airflow', } secrets_manager_backend = SecretsManagerBackend() secrets_manager_backend.client.put_secret_value(**param) assert secrets_manager_backend.get_conn_uri(conn_id=conn_id) is None assert [] == secrets_manager_backend.get_connections(conn_id=conn_id)
def test_aws_secrets_manager_get_connections(self, mock_get_uri): mock_get_uri.return_value = "scheme://*****:*****@host:100" conn_list = SecretsManagerBackend().get_connections("fake_conn") conn = conn_list[0] assert conn.host == 'host'