Skip to content

Commit 81b87d4

Browse files
rootcssrootcss
authored
Add unit tests for GcpBodyFieldSanitizer in Google providers (#9996)
1 parent a28c9c6 commit 81b87d4

File tree

3 files changed

+236
-3
lines changed

3 files changed

+236
-3
lines changed

β€Žairflow/providers/google/cloud/utils/field_sanitizer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,8 @@
6868
>>> }
6969
>>> }
7070
>>> sanitizer=GcpBodyFieldSanitizer(FIELDS_TO_SANITIZE)
71-
>>> SANITIZED_BODY = sanitizer.sanitize(body)
72-
>>> json.dumps(SANITIZED_BODY, indent=2)
71+
>>> sanitizer.sanitize(body)
72+
>>> json.dumps(body, indent=2)
7373
{
7474
"name": "instance",
7575
"properties": {

β€Žtests/providers/google/cloud/utils/test_field_sanitizer.py

Lines changed: 234 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,234 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
import unittest
19+
from copy import deepcopy
20+
21+
from airflow.providers.google.cloud.utils.field_sanitizer import GcpBodyFieldSanitizer
22+
23+
24+
class TestGcpBodyFieldSanitizer(unittest.TestCase):
25+
def test_sanitize_should_sanitize_empty_body_and_fields(self):
26+
body = {}
27+
fields_to_sanitize = []
28+
29+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
30+
sanitizer.sanitize(body)
31+
32+
self.assertEqual({}, body)
33+
34+
def test_sanitize_should_not_fail_with_none_body(self):
35+
body = None
36+
fields_to_sanitize = []
37+
38+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
39+
sanitizer.sanitize(body)
40+
41+
self.assertIsNone(body)
42+
43+
def test_sanitize_should_fail_with_none_fields(self):
44+
body = {}
45+
fields_to_sanitize = None
46+
47+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
48+
49+
with self.assertRaises(TypeError):
50+
sanitizer.sanitize(body)
51+
52+
def test_sanitize_should_not_fail_if_field_is_absent_in_body(self):
53+
body = {}
54+
fields_to_sanitize = ["kind"]
55+
56+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
57+
sanitizer.sanitize(body)
58+
59+
self.assertEqual({}, body)
60+
61+
def test_sanitize_should_not_remove_fields_for_incorrect_specification(self):
62+
actual_body = [
63+
{"kind": "compute#instanceTemplate", "name": "instance"},
64+
{"kind": "compute#instanceTemplate1", "name": "instance1"},
65+
{"kind": "compute#instanceTemplate2", "name": "instance2"},
66+
]
67+
body = deepcopy(actual_body)
68+
fields_to_sanitize = ["kind"]
69+
70+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
71+
sanitizer.sanitize(body)
72+
73+
self.assertEqual(actual_body, body)
74+
75+
def test_sanitize_should_remove_all_fields_from_root_level(self):
76+
body = {"kind": "compute#instanceTemplate", "name": "instance"}
77+
fields_to_sanitize = ["kind"]
78+
79+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
80+
sanitizer.sanitize(body)
81+
82+
self.assertEqual({"name": "instance"}, body)
83+
84+
def test_sanitize_should_remove_for_multiple_fields_from_root_level(self):
85+
body = {"kind": "compute#instanceTemplate", "name": "instance"}
86+
fields_to_sanitize = ["kind", "name"]
87+
88+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
89+
sanitizer.sanitize(body)
90+
91+
self.assertEqual({}, body)
92+
93+
def test_sanitize_should_remove_all_fields_in_a_list_value(self):
94+
body = {"fields": [
95+
{"kind": "compute#instanceTemplate", "name": "instance"},
96+
{"kind": "compute#instanceTemplate1", "name": "instance1"},
97+
{"kind": "compute#instanceTemplate2", "name": "instance2"},
98+
]}
99+
fields_to_sanitize = ["fields.kind"]
100+
101+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
102+
sanitizer.sanitize(body)
103+
104+
self.assertEqual({"fields": [
105+
{"name": "instance"},
106+
{"name": "instance1"},
107+
{"name": "instance2"},
108+
]}, body)
109+
110+
def test_sanitize_should_remove_all_fields_in_any_nested_body(self):
111+
fields_to_sanitize = [
112+
"kind",
113+
"properties.disks.kind",
114+
"properties.metadata.kind",
115+
]
116+
117+
body = {
118+
"kind": "compute#instanceTemplate",
119+
"name": "instance",
120+
"properties": {
121+
"disks": [
122+
{
123+
"name": "a",
124+
"kind": "compute#attachedDisk",
125+
"type": "PERSISTENT",
126+
"mode": "READ_WRITE",
127+
},
128+
{
129+
"name": "b",
130+
"kind": "compute#attachedDisk",
131+
"type": "PERSISTENT",
132+
"mode": "READ_WRITE",
133+
}
134+
],
135+
"metadata": {
136+
"kind": "compute#metadata",
137+
"fingerprint": "GDPUYxlwHe4="
138+
},
139+
}
140+
}
141+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
142+
sanitizer.sanitize(body)
143+
144+
self.assertEqual({
145+
"name": "instance",
146+
"properties": {
147+
"disks": [
148+
{
149+
"name": "a",
150+
"type": "PERSISTENT",
151+
"mode": "READ_WRITE"
152+
},
153+
{
154+
"name": "b",
155+
"type": "PERSISTENT",
156+
"mode": "READ_WRITE"
157+
}
158+
],
159+
"metadata": {
160+
"fingerprint": "GDPUYxlwHe4="
161+
}
162+
}
163+
}, body)
164+
165+
def test_sanitize_should_not_fail_if_specification_has_none_value(self):
166+
fields_to_sanitize = [
167+
"kind",
168+
"properties.disks.kind",
169+
"properties.metadata.kind",
170+
]
171+
172+
body = {
173+
"kind": "compute#instanceTemplate",
174+
"name": "instance",
175+
"properties": {
176+
"disks": None
177+
}
178+
}
179+
180+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
181+
sanitizer.sanitize(body)
182+
183+
self.assertEqual({
184+
"name": "instance",
185+
"properties": {
186+
"disks": None
187+
}
188+
}, body)
189+
190+
def test_sanitize_should_not_fail_if_no_specification_matches(self):
191+
fields_to_sanitize = [
192+
"properties.disks.kind1",
193+
"properties.metadata.kind2",
194+
]
195+
196+
body = {
197+
"name": "instance",
198+
"properties": {
199+
"disks": None
200+
}
201+
}
202+
203+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
204+
sanitizer.sanitize(body)
205+
206+
self.assertEqual({
207+
"name": "instance",
208+
"properties": {
209+
"disks": None
210+
}
211+
}, body)
212+
213+
def test_sanitize_should_not_fail_if_type_in_body_do_not_match_with_specification(self):
214+
fields_to_sanitize = [
215+
"properties.disks.kind",
216+
"properties.metadata.kind2",
217+
]
218+
219+
body = {
220+
"name": "instance",
221+
"properties": {
222+
"disks": 1
223+
}
224+
}
225+
226+
sanitizer = GcpBodyFieldSanitizer(fields_to_sanitize)
227+
sanitizer.sanitize(body)
228+
229+
self.assertEqual({
230+
"name": "instance",
231+
"properties": {
232+
"disks": 1
233+
}
234+
}, body)

β€Žtests/test_project_structure.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@
3030
MISSING_TEST_FILES = {
3131
'tests/providers/google/cloud/log/test_gcs_task_handler.py',
3232
'tests/providers/google/cloud/operators/test_datastore.py',
33-
'tests/providers/google/cloud/utils/test_field_sanitizer.py',
3433
'tests/providers/google/cloud/utils/test_mlengine_prediction_summary.py',
3534
'tests/providers/microsoft/azure/sensors/test_azure_cosmos.py',
3635
'tests/providers/microsoft/azure/log/test_wasb_task_handler.py',

0 commit comments

Comments
 (0)