-
Notifications
You must be signed in to change notification settings - Fork 13.7k
/
test_secrets_backends.py
102 lines (88 loc) · 4 KB
/
test_secrets_backends.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations
import os
from unittest import mock
import pytest
from airflow.models.connection import Connection
from airflow.models.variable import Variable
from airflow.secrets.base_secrets import BaseSecretsBackend
from airflow.secrets.environment_variables import EnvironmentVariablesBackend
from airflow.secrets.metastore import MetastoreBackend
from airflow.utils.session import create_session
from tests.test_utils.db import clear_db_connections, clear_db_variables
class SampleConn:
def __init__(self, conn_id, variation: str):
self.conn_id = conn_id
self.var_name = "AIRFLOW_CONN_" + self.conn_id.upper()
self.host = f"host_{variation}.com"
self.conn_uri = "mysql://user:pw@" + self.host + "/schema?extra1=val%2B1&extra2=val%2B2"
self.conn = Connection(conn_id=self.conn_id, uri=self.conn_uri)
class TestBaseSecretsBackend:
def setup_method(self) -> None:
clear_db_variables()
def teardown_method(self) -> None:
clear_db_connections()
clear_db_variables()
@pytest.mark.parametrize(
"kwargs, output",
[
({"path_prefix": "PREFIX", "secret_id": "ID"}, "PREFIX/ID"),
({"path_prefix": "PREFIX", "secret_id": "ID", "sep": "-"}, "PREFIX-ID"),
],
ids=["default", "with_sep"],
)
def test_build_path(self, kwargs, output):
build_path = BaseSecretsBackend.build_path
assert build_path(**kwargs) == output
def test_connection_env_secrets_backend(self):
sample_conn_1 = SampleConn("sample_1", "A")
env_secrets_backend = EnvironmentVariablesBackend()
os.environ[sample_conn_1.var_name] = sample_conn_1.conn_uri
conn = env_secrets_backend.get_connection(sample_conn_1.conn_id)
# we could make this more precise by defining __eq__ method for Connection
assert sample_conn_1.host.lower() == conn.host
def test_connection_metastore_secrets_backend(self):
sample_conn_2 = SampleConn("sample_2", "A")
with create_session() as session:
session.add(sample_conn_2.conn)
session.commit()
metastore_backend = MetastoreBackend()
conn = metastore_backend.get_connection("sample_2")
assert sample_conn_2.host.lower() == conn.host
@mock.patch.dict(
'os.environ',
{
'AIRFLOW_VAR_HELLO': 'World',
'AIRFLOW_VAR_EMPTY_STR': '',
},
)
def test_variable_env_secrets_backend(self):
env_secrets_backend = EnvironmentVariablesBackend()
variable_value = env_secrets_backend.get_variable(key="hello")
assert 'World' == variable_value
assert env_secrets_backend.get_variable(key="non_existent_key") is None
assert '' == env_secrets_backend.get_variable(key="empty_str")
def test_variable_metastore_secrets_backend(self):
Variable.set(key="hello", value="World")
Variable.set(key="empty_str", value="")
metastore_backend = MetastoreBackend()
variable_value = metastore_backend.get_variable(key="hello")
assert "World" == variable_value
assert metastore_backend.get_variable(key="non_existent_key") is None
assert '' == metastore_backend.get_variable(key="empty_str")