Fix mock_get_connection in test_redis

apache · Dec 5, 2024 · 4ae21b0 · 4ae21b0
1 parent 8b69bf8
commit 4ae21b0
Showing 1 changed file with 27 additions and 25 deletions.
diff --git a/providers/tests/redis/hooks/test_redis.py b/providers/tests/redis/hooks/test_redis.py
@@ -41,39 +41,41 @@ def test_get_conn(self):
     @mock.patch("airflow.providers.redis.hooks.redis.Redis")
     @mock.patch(
         "airflow.providers.redis.hooks.redis.RedisHook.get_connection",
-        return_value=Connection(
+    )
+    def test_get_conn_with_extra_config(self, mock_get_connection, mock_redis):
+        connection = Connection(
             login="user",
             password="password",
             host="remote_host",
             port=1234,
-            extra="""{
-                        "db": 2,
-                        "ssl": true,
-                        "ssl_cert_reqs": "required",
-                        "ssl_ca_certs": "/path/to/custom/ca-cert",
-                        "ssl_keyfile": "/path/to/key-file",
-                        "ssl_certfile": "/path/to/cert-file",
-                        "ssl_check_hostname": true
-                    }""",
-        ),
-    )
-    def test_get_conn_with_extra_config(self, mock_get_connection, mock_redis):
-        connection = mock_get_connection.return_value
+        )
+        connection.set_extra(
+            """{
+                "db": 2,
+                "ssl": true,
+                "ssl_cert_reqs": "required",
+                "ssl_ca_certs": "/path/to/custom/ca-cert",
+                "ssl_keyfile": "/path/to/key-file",
+                "ssl_certfile": "/path/to/cert-file",
+                "ssl_check_hostname": true
+            }"""
+        )
+        mock_get_connection.return_value = connection
         hook = RedisHook()
 
         hook.get_conn()
         mock_redis.assert_called_once_with(
-            host=connection.host,
-            username=connection.login,
-            password=connection.password,
-            port=connection.port,
-            db=connection.extra_dejson["db"],
-            ssl=connection.extra_dejson["ssl"],
-            ssl_cert_reqs=connection.extra_dejson["ssl_cert_reqs"],
-            ssl_ca_certs=connection.extra_dejson["ssl_ca_certs"],
-            ssl_keyfile=connection.extra_dejson["ssl_keyfile"],
-            ssl_certfile=connection.extra_dejson["ssl_certfile"],
-            ssl_check_hostname=connection.extra_dejson["ssl_check_hostname"],
+            host=mock_get_connection.host,
+            username=mock_get_connection.login,
+            password=mock_get_connection.password,
+            port=mock_get_connection.port,
+            db=mock_get_connection.extra_dejson["db"],
+            ssl=mock_get_connection.extra_dejson["ssl"],
+            ssl_cert_reqs=mock_get_connection.extra_dejson["ssl_cert_reqs"],
+            ssl_ca_certs=mock_get_connection.extra_dejson["ssl_ca_certs"],
+            ssl_keyfile=mock_get_connection.extra_dejson["ssl_keyfile"],
+            ssl_certfile=mock_get_connection.extra_dejson["ssl_certfile"],
+            ssl_check_hostname=mock_get_connection.extra_dejson["ssl_check_hostname"],
         )
 
     def test_get_conn_password_stays_none(self):