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

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 

22import pytest 

23 

24from airflow.configuration import ensure_secrets_loaded, initialize_secrets_backends 

25from airflow.models import Connection, Variable 

26from airflow.secrets.cache import SecretCache 

27from tests.test_utils.config import conf_vars 

28from tests.test_utils.db import clear_db_variables 

29 

30 

31class TestConnectionsFromSecrets: 

32 def setup_method(self) -> None: 

33 SecretCache.reset() 

34 

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

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

37 def test_get_connection_second_try(self, mock_env_get, mock_meta_get): 

38 mock_env_get.side_effect = [None] # return None 

39 Connection.get_connection_from_secrets("fake_conn_id") 

40 mock_meta_get.assert_called_once_with(conn_id="fake_conn_id") 

41 mock_env_get.assert_called_once_with(conn_id="fake_conn_id") 

42 

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

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

45 def test_get_connection_first_try(self, mock_env_get, mock_meta_get): 

46 mock_env_get.return_value = Connection("something") # returns something 

47 Connection.get_connection_from_secrets("fake_conn_id") 

48 mock_env_get.assert_called_once_with(conn_id="fake_conn_id") 

49 mock_meta_get.assert_not_called() 

50 

51 @conf_vars( 

52 { 

53 ( 

54 "secrets", 

55 "backend", 

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

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

58 } 

59 ) 

60 def test_initialize_secrets_backends(self): 

61 backends = initialize_secrets_backends() 

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

63 

64 assert 3 == len(backends) 

65 assert "SystemsManagerParameterStoreBackend" in backend_classes 

66 

67 @conf_vars( 

68 { 

69 ( 

70 "secrets", 

71 "backend", 

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

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

74 } 

75 ) 

76 def test_backends_kwargs(self): 

77 backends = initialize_secrets_backends() 

78 systems_manager = next( 

79 backend 

80 for backend in backends 

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

82 ) 

83 assert systems_manager.kwargs == {} 

84 assert systems_manager.use_ssl is False 

85 

86 @conf_vars( 

87 { 

88 ( 

89 "secrets", 

90 "backend", 

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

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

93 } 

94 ) 

95 @mock.patch.dict( 

96 "os.environ", 

97 { 

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

99 }, 

100 ) 

101 @mock.patch( 

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

103 "SystemsManagerParameterStoreBackend.get_connection" 

104 ) 

105 def test_backend_fallback_to_env_var(self, mock_get_connection): 

106 mock_get_connection.return_value = None 

107 

108 backends = ensure_secrets_loaded() 

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

110 assert "SystemsManagerParameterStoreBackend" in backend_classes 

111 

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

113 

114 # Assert that SystemsManagerParameterStoreBackend.get_conn_uri was called 

115 mock_get_connection.assert_called_once_with(conn_id="test_mysql") 

116 

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

118 

119 

120@pytest.mark.db_test 

121class TestVariableFromSecrets: 

122 def setup_method(self) -> None: 

123 clear_db_variables() 

124 SecretCache.reset() 

125 

126 def teardown_method(self) -> None: 

127 clear_db_variables() 

128 

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

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

131 def test_get_variable_second_try(self, mock_env_get, mock_meta_get): 

132 """ 

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

134 Metastore DB 

135 """ 

136 mock_env_get.return_value = None 

137 mock_meta_get.return_value = "val" 

138 

139 Variable.get_variable_from_secrets("fake_var_key") 

140 

141 mock_meta_get.assert_called_once_with(key="fake_var_key") 

142 mock_env_get.assert_called_once_with(key="fake_var_key") 

143 

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

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

146 def test_get_variable_first_try(self, mock_env_get, mock_meta_get): 

147 """ 

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

149 Metastore DB 

150 """ 

151 mock_env_get.return_value = "something" 

152 Variable.get_variable_from_secrets("fake_var_key") 

153 mock_env_get.assert_called_once_with(key="fake_var_key") 

154 mock_meta_get.assert_not_called() 

155 

156 def test_backend_fallback_to_default_var(self): 

157 """ 

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

159 the value returned is default_var 

160 """ 

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

162 assert "new" == variable_value 

163 

164 @conf_vars( 

165 { 

166 ( 

167 "secrets", 

168 "backend", 

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

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

171 } 

172 ) 

173 @mock.patch.dict( 

174 "os.environ", 

175 { 

176 "AIRFLOW_VAR_MYVAR": "a_venv_value", 

177 }, 

178 ) 

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

180 @mock.patch( 

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

182 "SystemsManagerParameterStoreBackend.get_variable" 

183 ) 

184 def test_backend_variable_order(self, mock_secret_get, mock_meta_get): 

185 backends = ensure_secrets_loaded() 

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

187 assert "SystemsManagerParameterStoreBackend" in backend_classes 

188 

189 mock_secret_get.return_value = None 

190 mock_meta_get.return_value = None 

191 

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

193 mock_secret_get.assert_called_with(key="MYVAR") 

194 mock_meta_get.assert_not_called() 

195 

196 mock_secret_get.return_value = None 

197 mock_meta_get.return_value = "a_metastore_value" 

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

199 mock_meta_get.assert_called_once_with(key="not_myvar") 

200 

201 mock_secret_get.return_value = "a_secret_value" 

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