Add unit tests for GcpBodyFieldSanitizer in Google providers (#9996)

apache · Jul 26, 2020 · 81b87d4 · 81b87d4
1 parent a28c9c6
commit 81b87d4
Show file tree

Hide file tree

Showing 3 changed files with 236 additions and 3 deletions.
diff --git a/airflow/providers/google/cloud/utils/field_sanitizer.py b/airflow/providers/google/cloud/utils/field_sanitizer.py
@@ -68,8 +68,8 @@
 >>> }
 >>> }
 >>> sanitizer=GcpBodyFieldSanitizer(FIELDS_TO_SANITIZE)
->>> SANITIZED_BODY = sanitizer.sanitize(body)
->>> json.dumps(SANITIZED_BODY, indent=2)
+>>> sanitizer.sanitize(body)
+>>> json.dumps(body, indent=2)
 {
  "name": "instance",
  "properties": {

diff --git a/tests/providers/google/cloud/utils/test_field_sanitizer.py b/tests/providers/google/cloud/utils/test_field_sanitizer.py
@@ -0,0 +1,234 @@
+# 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.
+
+import unittest
+from copy import deepcopy
+
+from airflow.providers.google.cloud.utils.field_sanitizer import GcpBodyFieldSanitizer
+
+
+class TestGcpBodyFieldSanitizer(unittest.TestCase):
+ def test_sanitize_should_sanitize_empty_body_and_fields(self):
+ body = {}
+ fields_to_sanitize = []
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({}, body)
+
+ def test_sanitize_should_not_fail_with_none_body(self):
+ body = None
+ fields_to_sanitize = []
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertIsNone(body)
+
+ def test_sanitize_should_fail_with_none_fields(self):
+ body = {}
+ fields_to_sanitize = None
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+
+ with self.assertRaises(TypeError):
+ sanitizer.sanitize(body)
+
+ def test_sanitize_should_not_fail_if_field_is_absent_in_body(self):
+ body = {}
+ fields_to_sanitize = ["kind"]
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({}, body)
+
+ def test_sanitize_should_not_remove_fields_for_incorrect_specification(self):
+ actual_body = [
+ {"kind": "compute#instanceTemplate", "name": "instance"},
+ {"kind": "compute#instanceTemplate1", "name": "instance1"},
+ {"kind": "compute#instanceTemplate2", "name": "instance2"},
+ ]
+ body = deepcopy(actual_body)
+ fields_to_sanitize = ["kind"]
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual(actual_body, body)
+
+ def test_sanitize_should_remove_all_fields_from_root_level(self):
+ body = {"kind": "compute#instanceTemplate", "name": "instance"}
+ fields_to_sanitize = ["kind"]
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({"name": "instance"}, body)
+
+ def test_sanitize_should_remove_for_multiple_fields_from_root_level(self):
+ body = {"kind": "compute#instanceTemplate", "name": "instance"}
+ fields_to_sanitize = ["kind", "name"]
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({}, body)
+
+ def test_sanitize_should_remove_all_fields_in_a_list_value(self):
+ body = {"fields": [
+ {"kind": "compute#instanceTemplate", "name": "instance"},
+ {"kind": "compute#instanceTemplate1", "name": "instance1"},
+ {"kind": "compute#instanceTemplate2", "name": "instance2"},
+ ]}
+ fields_to_sanitize = ["fields.kind"]
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({"fields": [
+ {"name": "instance"},
+ {"name": "instance1"},
+ {"name": "instance2"},
+ ]}, body)
+
+ def test_sanitize_should_remove_all_fields_in_any_nested_body(self):
+ fields_to_sanitize = [
+ "kind",
+ "properties.disks.kind",
+ "properties.metadata.kind",
+ ]
+
+ body = {
+ "kind": "compute#instanceTemplate",
+ "name": "instance",
+ "properties": {
+ "disks": [
+ {
+ "name": "a",
+ "kind": "compute#attachedDisk",
+ "type": "PERSISTENT",
+ "mode": "READ_WRITE",
+ },
+ {
+ "name": "b",
+ "kind": "compute#attachedDisk",
+ "type": "PERSISTENT",
+ "mode": "READ_WRITE",
+ }
+ ],
+ "metadata": {
+ "kind": "compute#metadata",
+ "fingerprint": "GDPUYxlwHe4="
+ },
+ }
+ }
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({
+ "name": "instance",
+ "properties": {
+ "disks": [
+ {
+ "name": "a",
+ "type": "PERSISTENT",
+ "mode": "READ_WRITE"
+ },
+ {
+ "name": "b",
+ "type": "PERSISTENT",
+ "mode": "READ_WRITE"
+ }
+ ],
+ "metadata": {
+ "fingerprint": "GDPUYxlwHe4="
+ }
+ }
+ }, body)
+
+ def test_sanitize_should_not_fail_if_specification_has_none_value(self):
+ fields_to_sanitize = [
+ "kind",
+ "properties.disks.kind",
+ "properties.metadata.kind",
+ ]
+
+ body = {
+ "kind": "compute#instanceTemplate",
+ "name": "instance",
+ "properties": {
+ "disks": None
+ }
+ }
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({
+ "name": "instance",
+ "properties": {
+ "disks": None
+ }
+ }, body)
+
+ def test_sanitize_should_not_fail_if_no_specification_matches(self):
+ fields_to_sanitize = [
+ "properties.disks.kind1",
+ "properties.metadata.kind2",
+ ]
+
+ body = {
+ "name": "instance",
+ "properties": {
+ "disks": None
+ }
+ }
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({
+ "name": "instance",
+ "properties": {
+ "disks": None
+ }
+ }, body)
+
+ def test_sanitize_should_not_fail_if_type_in_body_do_not_match_with_specification(self):
+ fields_to_sanitize = [
+ "properties.disks.kind",
+ "properties.metadata.kind2",
+ ]
+
+ body = {
+ "name": "instance",
+ "properties": {
+ "disks": 1
+ }
+ }
+
+ sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
+ sanitizer.sanitize(body)
+
+ self.assertEqual({
+ "name": "instance",
+ "properties": {
+ "disks": 1
+ }
+ }, body)
diff --git a/tests/test_project_structure.py b/tests/test_project_structure.py
@@ -30,7 +30,6 @@
 MISSING_TEST_FILES = {
  'tests/providers/google/cloud/log/test_gcs_task_handler.py',
  'tests/providers/google/cloud/operators/test_datastore.py',
- 'tests/providers/google/cloud/utils/test_field_sanitizer.py',
  'tests/providers/google/cloud/utils/test_mlengine_prediction_summary.py',
  'tests/providers/microsoft/azure/sensors/test_azure_cosmos.py',
  'tests/providers/microsoft/azure/log/test_wasb_task_handler.py',