Add final set of protobuf helpers to api_core (#4259)

1 files changed, 165 insertions, 0 deletions

diff --git a/tests/unit/test_protobuf_helpers.py b/tests/unit/test_protobuf_helpers.py
index b9aca76..8f86aa4 100644
--- a/tests/unit/test_protobuf_helpers.py
+++ b/tests/unit/test_protobuf_helpers.py
@@ -14,8 +14,11 @@
 
 import pytest
 
+from google.api import http_pb2
 from google.api_core import protobuf_helpers
+from google.longrunning import operations_pb2
 from google.protobuf import any_pb2
+from google.protobuf import timestamp_pb2
 from google.protobuf.message import Message
 from google.type import date_pb2
 from google.type import timeofday_pb2
@@ -65,3 +68,165 @@ def test_get_messages():
     # Ensure that no non-Message objects were exported.
     for value in answer.values():
         assert issubclass(value, Message)
+
+
+def test_get_dict_absent():
+    with pytest.raises(KeyError):
+        assert protobuf_helpers.get({}, 'foo')
+
+
+def test_get_dict_present():
+    assert protobuf_helpers.get({'foo': 'bar'}, 'foo') == 'bar'
+
+
+def test_get_dict_default():
+    assert protobuf_helpers.get({}, 'foo', default='bar') == 'bar'
+
+
+def test_get_dict_nested():
+    assert protobuf_helpers.get({'foo': {'bar': 'baz'}}, 'foo.bar') == 'baz'
+
+
+def test_get_dict_nested_default():
+    assert protobuf_helpers.get({}, 'foo.baz', default='bacon') == 'bacon'
+    assert (
+        protobuf_helpers.get({'foo': {}}, 'foo.baz', default='bacon') ==
+        'bacon')
+
+
+def test_get_msg_sentinel():
+    msg = timestamp_pb2.Timestamp()
+    with pytest.raises(KeyError):
+        assert protobuf_helpers.get(msg, 'foo')
+
+
+def test_get_msg_present():
+    msg = timestamp_pb2.Timestamp(seconds=42)
+    assert protobuf_helpers.get(msg, 'seconds') == 42
+
+
+def test_get_msg_default():
+    msg = timestamp_pb2.Timestamp()
+    assert protobuf_helpers.get(msg, 'foo', default='bar') == 'bar'
+
+
+def test_invalid_object():
+    with pytest.raises(TypeError):
+        protobuf_helpers.get(object(), 'foo', 'bar')
+
+
+def test_set_dict():
+    mapping = {}
+    protobuf_helpers.set(mapping, 'foo', 'bar')
+    assert mapping == {'foo': 'bar'}
+
+
+def test_set_msg():
+    msg = timestamp_pb2.Timestamp()
+    protobuf_helpers.set(msg, 'seconds', 42)
+    assert msg.seconds == 42
+
+
+def test_set_dict_nested():
+    mapping = {}
+    protobuf_helpers.set(mapping, 'foo.bar', 'baz')
+    assert mapping == {'foo': {'bar': 'baz'}}
+
+
+def test_set_invalid_object():
+    with pytest.raises(TypeError):
+        protobuf_helpers.set(object(), 'foo', 'bar')
+
+
+def test_set_list():
+    list_ops_response = operations_pb2.ListOperationsResponse()
+
+    protobuf_helpers.set(list_ops_response, 'operations', [
+        {'name': 'foo'},
+        operations_pb2.Operation(name='bar'),
+    ])
+
+    assert len(list_ops_response.operations) == 2
+
+    for operation in list_ops_response.operations:
+        assert isinstance(operation, operations_pb2.Operation)
+
+    assert list_ops_response.operations[0].name == 'foo'
+    assert list_ops_response.operations[1].name == 'bar'
+
+
+def test_set_list_clear_existing():
+    list_ops_response = operations_pb2.ListOperationsResponse(
+        operations=[{'name': 'baz'}],
+    )
+
+    protobuf_helpers.set(list_ops_response, 'operations', [
+        {'name': 'foo'},
+        operations_pb2.Operation(name='bar'),
+    ])
+
+    assert len(list_ops_response.operations) == 2
+    for operation in list_ops_response.operations:
+        assert isinstance(operation, operations_pb2.Operation)
+    assert list_ops_response.operations[0].name == 'foo'
+    assert list_ops_response.operations[1].name == 'bar'
+
+
+def test_set_msg_with_msg_field():
+    rule = http_pb2.HttpRule()
+    pattern = http_pb2.CustomHttpPattern(kind='foo', path='bar')
+
+    protobuf_helpers.set(rule, 'custom', pattern)
+
+    assert rule.custom.kind == 'foo'
+    assert rule.custom.path == 'bar'
+
+
+def test_set_msg_with_dict_field():
+    rule = http_pb2.HttpRule()
+    pattern = {'kind': 'foo', 'path': 'bar'}
+
+    protobuf_helpers.set(rule, 'custom', pattern)
+
+    assert rule.custom.kind == 'foo'
+    assert rule.custom.path == 'bar'
+
+
+def test_set_msg_nested_key():
+    rule = http_pb2.HttpRule(
+        custom=http_pb2.CustomHttpPattern(kind='foo', path='bar'))
+
+    protobuf_helpers.set(rule, 'custom.kind', 'baz')
+
+    assert rule.custom.kind == 'baz'
+    assert rule.custom.path == 'bar'
+
+
+def test_setdefault_dict_unset():
+    mapping = {}
+    protobuf_helpers.setdefault(mapping, 'foo', 'bar')
+    assert mapping == {'foo': 'bar'}
+
+
+def test_setdefault_dict_falsy():
+    mapping = {'foo': None}
+    protobuf_helpers.setdefault(mapping, 'foo', 'bar')
+    assert mapping == {'foo': 'bar'}
+
+
+def test_setdefault_dict_truthy():
+    mapping = {'foo': 'bar'}
+    protobuf_helpers.setdefault(mapping, 'foo', 'baz')
+    assert mapping == {'foo': 'bar'}
+
+
+def test_setdefault_pb2_falsy():
+    operation = operations_pb2.Operation()
+    protobuf_helpers.setdefault(operation, 'name', 'foo')
+    assert operation.name == 'foo'
+
+
+def test_setdefault_pb2_truthy():
+    operation = operations_pb2.Operation(name='bar')
+    protobuf_helpers.setdefault(operation, 'name', 'foo')
+    assert operation.name == 'bar'