Google Memcached hooks - improve protobuf messages handling (#11743)

apache · Oct 26, 2020 · 240c7d4 · 240c7d4
1 parent 0390515
commit 240c7d4
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 26 deletions.
diff --git a/airflow/providers/google/cloud/hooks/cloud_memorystore.py b/airflow/providers/google/cloud/hooks/cloud_memorystore.py
@@ -17,7 +17,6 @@
 # under the License.
 """Hooks for Cloud Memorystore service"""
 from typing import Dict, Optional, Sequence, Tuple, Union
-import json
 
 from google.api_core.exceptions import NotFound
 from google.api_core import path_template
@@ -28,7 +27,6 @@
 from google.cloud.redis_v1.gapic.enums import FailoverInstanceRequest
 from google.cloud.redis_v1.types import FieldMask, InputConfig, Instance, OutputConfig
 from google.protobuf.json_format import ParseDict
-import proto
 
 from airflow import version
 from airflow.exceptions import AirflowException
@@ -569,11 +567,6 @@ def _append_label(instance: cloud_memcache.Instance, key: str, val: str) -> clou
         instance.labels.update({key: val})
         return instance
 
-    @staticmethod
-    def proto_message_to_dict(message: proto.Message) -> dict:
-        """Helper method to parse protobuf message to dictionary."""
-        return json.loads(message.__class__.to_json(message))
-
     @GoogleBaseHook.fallback_to_default_project_id
     def apply_parameters(
         self,

diff --git a/airflow/providers/google/cloud/operators/cloud_memorystore.py b/airflow/providers/google/cloud/operators/cloud_memorystore.py
@@ -1290,7 +1290,7 @@ def execute(self, context: Dict):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        return hook.proto_message_to_dict(result)
+        return cloud_memcache.Instance.to_dict(result)
 
 
 class CloudMemorystoreMemcachedDeleteInstanceOperator(BaseOperator):
@@ -1438,7 +1438,7 @@ def execute(self, context: Dict):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        return hook.proto_message_to_dict(result)
+        return cloud_memcache.Instance.to_dict(result)
 
 
 class CloudMemorystoreMemcachedListInstancesOperator(BaseOperator):
@@ -1520,7 +1520,7 @@ def execute(self, context: Dict):
             timeout=self.timeout,
             metadata=self.metadata,
         )
-        instances = [hook.proto_message_to_dict(a) for a in result]
+        instances = [cloud_memcache.Instance.to_dict(a) for a in result]
         return instances
 
 

diff --git a/tests/providers/google/cloud/hooks/test_cloud_memorystore.py b/tests/providers/google/cloud/hooks/test_cloud_memorystore.py
@@ -599,19 +599,3 @@ def test_update_instance(self, mock_get_conn, mock_project_id):
             timeout=TEST_TIMEOUT,
             metadata=TEST_METADATA,
         )
-
-    def test_proto_functions(self):
-        instance_dict = {
-            'name': 'test_name',
-            'node_count': 1,
-            'node_config': {'cpu_count': 1, 'memory_size_mb': 1024},
-        }
-        instance = cloud_memcache.Instance(instance_dict)
-        instance_dict_result = self.hook.proto_message_to_dict(instance)
-        self.assertEqual(instance_dict_result["name"], instance_dict["name"])
-        self.assertEqual(
-            instance_dict_result["nodeConfig"]["cpuCount"], instance_dict["node_config"]["cpu_count"]
-        )
-        self.assertEqual(
-            instance_dict_result["nodeConfig"]["memorySizeMb"], instance_dict["node_config"]["memory_size_mb"]
-        )
diff --git a/tests/providers/google/cloud/operators/test_cloud_memorystore.py b/tests/providers/google/cloud/operators/test_cloud_memorystore.py
@@ -21,6 +21,7 @@
 from google.api_core.retry import Retry
 from google.cloud.redis_v1.gapic.enums import FailoverInstanceRequest
 from google.cloud.redis_v1.types import Instance
+from google.cloud.memcache_v1beta2.types import cloud_memcache
 
 from airflow.providers.google.cloud.operators.cloud_memorystore import (
     CloudMemorystoreCreateInstanceAndImportOperator,
@@ -386,6 +387,7 @@ def test_assert_valid_hook_call(self, mock_hook):
 class TestCloudMemorystoreMemcachedCreateInstanceOperator(TestCase):
     @mock.patch("airflow.providers.google.cloud.operators.cloud_memorystore.CloudMemorystoreMemcachedHook")
     def test_assert_valid_hook_call(self, mock_hook):
+        mock_hook.return_value.create_instance.return_value = cloud_memcache.Instance()
         task = CloudMemorystoreMemcachedCreateInstanceOperator(
             task_id=TEST_TASK_ID,
             location=TEST_LOCATION,
@@ -438,6 +440,7 @@ def test_assert_valid_hook_call(self, mock_hook):
 class TestCloudMemorystoreMemcachedGetInstanceOperator(TestCase):
     @mock.patch("airflow.providers.google.cloud.operators.cloud_memorystore.CloudMemorystoreMemcachedHook")
     def test_assert_valid_hook_call(self, mock_hook):
+        mock_hook.return_value.get_instance.return_value = cloud_memcache.Instance()
         task = CloudMemorystoreMemcachedGetInstanceOperator(
             task_id=TEST_TASK_ID,
             location=TEST_LOCATION,