Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/tests/always/test_secrets.py: 9%

1# 

2# Licensed to the Apache Software Foundation (ASF) under one 

3# or more contributor license agreements. See the NOTICE file 

4# distributed with this work for additional information 

5# regarding copyright ownership. The ASF licenses this file 

6# to you under the Apache License, Version 2.0 (the 

7# "License"); you may not use this file except in compliance 

8# with the License. You may obtain a copy of the License at 

9# 

10# http://www.apache.org/licenses/LICENSE-2.0 

11# 

12# Unless required by applicable law or agreed to in writing, 

13# software distributed under the License is distributed on an 

14# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

15# KIND, either express or implied. See the License for the 

16# specific language governing permissions and limitations 

17# under the License. 

18from __future__ import annotations 

19 

20from unittest import mock 

21 

22from airflow.configuration import ensure_secrets_loaded, initialize_secrets_backends 

23from airflow.models import Connection, Variable 

24from tests.test_utils.config import conf_vars 

25from tests.test_utils.db import clear_db_variables 

26 

27 

28class TestConnectionsFromSecrets: 

29 @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_connection") 

30 @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_connection") 

31 def test_get_connection_second_try(self, mock_env_get, mock_meta_get): 

32 mock_env_get.side_effect = [None] # return None 

33 Connection.get_connection_from_secrets("fake_conn_id") 

34 mock_meta_get.assert_called_once_with(conn_id="fake_conn_id") 

35 mock_env_get.assert_called_once_with(conn_id="fake_conn_id") 

36 

37 @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_connection") 

38 @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_connection") 

39 def test_get_connection_first_try(self, mock_env_get, mock_meta_get): 

40 mock_env_get.side_effect = ["something"] # returns something 

41 Connection.get_connection_from_secrets("fake_conn_id") 

42 mock_env_get.assert_called_once_with(conn_id="fake_conn_id") 

43 mock_meta_get.not_called() 

44 

45 @conf_vars( 

46 { 

47 ( 

48 "secrets", 

49 "backend", 

50 ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", 

51 ("secrets", "backend_kwargs"): '{"connections_prefix": "/airflow", "profile_name": null}', 

52 } 

53 ) 

54 def test_initialize_secrets_backends(self): 

55 backends = initialize_secrets_backends() 

56 backend_classes = [backend.__class__.__name__ for backend in backends] 

57 

58 assert 3 == len(backends) 

59 assert "SystemsManagerParameterStoreBackend" in backend_classes 

60 

61 @conf_vars( 

62 { 

63 ( 

64 "secrets", 

65 "backend", 

66 ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", 

67 ("secrets", "backend_kwargs"): '{"use_ssl": false}', 

68 } 

69 ) 

70 def test_backends_kwargs(self): 

71 backends = initialize_secrets_backends() 

72 systems_manager = next( 

73 backend 

74 for backend in backends 

75 if backend.__class__.__name__ == "SystemsManagerParameterStoreBackend" 

76 ) 

77 assert systems_manager.kwargs == {} 

78 assert systems_manager.use_ssl is False 

79 

80 @conf_vars( 

81 { 

82 ( 

83 "secrets", 

84 "backend", 

85 ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", 

86 ("secrets", "backend_kwargs"): '{"connections_prefix": "/airflow", "profile_name": null}', 

87 } 

88 ) 

89 @mock.patch.dict( 

90 "os.environ", 

91 { 

92 "AIRFLOW_CONN_TEST_MYSQL": "mysql://airflow:airflow@host:5432/airflow", 

93 }, 

94 ) 

95 @mock.patch( 

96 "airflow.providers.amazon.aws.secrets.systems_manager." 

97 "SystemsManagerParameterStoreBackend.get_connection" 

98 ) 

99 def test_backend_fallback_to_env_var(self, mock_get_connection): 

100 mock_get_connection.return_value = None 

101 

102 backends = ensure_secrets_loaded() 

103 backend_classes = [backend.__class__.__name__ for backend in backends] 

104 assert "SystemsManagerParameterStoreBackend" in backend_classes 

105 

106 conn = Connection.get_connection_from_secrets(conn_id="test_mysql") 

107 

108 # Assert that SystemsManagerParameterStoreBackend.get_conn_uri was called 

109 mock_get_connection.assert_called_once_with(conn_id="test_mysql") 

110 

111 assert "mysql://airflow:airflow@host:5432/airflow" == conn.get_uri() 

112 

113 

114class TestVariableFromSecrets: 

115 def setup_method(self) -> None: 

116 clear_db_variables() 

117 

118 def teardown_method(self) -> None: 

119 clear_db_variables() 

120 

121 @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_variable") 

122 @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_variable") 

123 def test_get_variable_second_try(self, mock_env_get, mock_meta_get): 

124 """ 

125 Test if Variable is not present in Environment Variable, it then looks for it in 

126 Metastore DB 

127 """ 

128 mock_env_get.return_value = None 

129 Variable.get_variable_from_secrets("fake_var_key") 

130 mock_meta_get.assert_called_once_with(key="fake_var_key") 

131 mock_env_get.assert_called_once_with(key="fake_var_key") 

132 

133 @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_variable") 

134 @mock.patch("airflow.secrets.environment_variables.EnvironmentVariablesBackend.get_variable") 

135 def test_get_variable_first_try(self, mock_env_get, mock_meta_get): 

136 """ 

137 Test if Variable is present in Environment Variable, it does not look for it in 

138 Metastore DB 

139 """ 

140 mock_env_get.return_value = "something" 

141 Variable.get_variable_from_secrets("fake_var_key") 

142 mock_env_get.assert_called_once_with(key="fake_var_key") 

143 mock_meta_get.assert_not_called() 

144 

145 def test_backend_fallback_to_default_var(self): 

146 """ 

147 Test if a default_var is defined and no backend has the Variable, 

148 the value returned is default_var 

149 """ 

150 variable_value = Variable.get(key="test_var", default_var="new") 

151 assert "new" == variable_value 

152 

153 @conf_vars( 

154 { 

155 ( 

156 "secrets", 

157 "backend", 

158 ): "airflow.providers.amazon.aws.secrets.systems_manager.SystemsManagerParameterStoreBackend", 

159 ("secrets", "backend_kwargs"): '{"variables_prefix": "/airflow", "profile_name": null}', 

160 } 

161 ) 

162 @mock.patch.dict( 

163 "os.environ", 

164 { 

165 "AIRFLOW_VAR_MYVAR": "a_venv_value", 

166 }, 

167 ) 

168 @mock.patch("airflow.secrets.metastore.MetastoreBackend.get_variable") 

169 @mock.patch( 

170 "airflow.providers.amazon.aws.secrets.systems_manager." 

171 "SystemsManagerParameterStoreBackend.get_variable" 

172 ) 

173 def test_backend_variable_order(self, mock_secret_get, mock_meta_get): 

174 backends = ensure_secrets_loaded() 

175 backend_classes = [backend.__class__.__name__ for backend in backends] 

176 assert "SystemsManagerParameterStoreBackend" in backend_classes 

177 

178 mock_secret_get.return_value = None 

179 mock_meta_get.return_value = None 

180 

181 assert "a_venv_value" == Variable.get(key="MYVAR") 

182 mock_secret_get.assert_called_with(key="MYVAR") 

183 mock_meta_get.assert_not_called() 

184 

185 mock_secret_get.return_value = None 

186 mock_meta_get.return_value = "a_metastore_value" 

187 assert "a_metastore_value" == Variable.get(key="not_myvar") 

188 mock_meta_get.assert_called_once_with(key="not_myvar") 

189 

190 mock_secret_get.return_value = "a_secret_value" 

191 assert "a_secret_value" == Variable.get(key="not_myvar")