From e3a1eba74dd8b67bcc73a78f784189ef2a9927c2 Mon Sep 17 00:00:00 2001
From: Daniel Sanche <sanche@google.com>
Date: Tue, 5 Oct 2021 16:25:27 -0700
Subject: [PATCH] feat: added support for iam AuditData proto (#396)

---
 google/cloud/logging_v2/entries.py |  1 +
 setup.py                           |  1 +
 tests/system/test_system.py        | 33 ++++++++++++++++++++++++++++++
 3 files changed, 35 insertions(+)

diff --git a/google/cloud/logging_v2/entries.py b/google/cloud/logging_v2/entries.py
index fa7e5d9d..0af5a46f 100644
--- a/google/cloud/logging_v2/entries.py
+++ b/google/cloud/logging_v2/entries.py
@@ -30,6 +30,7 @@
 # import officially supported proto definitions
 import google.cloud.audit.audit_log_pb2  # noqa: F401
 import google.cloud.appengine_logging  # noqa: F401
+from google.iam.v1.logging import audit_data_pb2  # noqa: F401
 
 _GLOBAL_RESOURCE = Resource(type="global", labels={})
 
diff --git a/setup.py b/setup.py
index 6a9ed8a5..5b882211 100644
--- a/setup.py
+++ b/setup.py
@@ -39,6 +39,7 @@
     # Until this issue is closed
     # https://github.com/googleapis/google-cloud-python/issues/10566
     "google-cloud-core >= 1.4.1, <3.0.0dev",
+    "grpc-google-iam-v1 >= 0.12.3, < 0.13dev",
     "proto-plus >= 1.11.0",
     "packaging >= 14.3",
 ]
diff --git a/tests/system/test_system.py b/tests/system/test_system.py
index 81de866e..532eea96 100644
--- a/tests/system/test_system.py
+++ b/tests/system/test_system.py
@@ -260,6 +260,39 @@ def test_list_entry_with_requestlog(self):
             protobuf_entry.to_api_repr()["protoPayload"]["@type"], type_url
         )
 
+    def test_list_entry_with_auditdata(self):
+        """
+        Test emitting and listing logs containing a google.iam.v1.logging.AuditData proto message
+        """
+        from google.protobuf import descriptor_pool
+        from google.cloud.logging_v2 import entries
+
+        pool = descriptor_pool.Default()
+        type_name = "google.iam.v1.logging.AuditData"
+        type_url = "type.googleapis.com/" + type_name
+        # Make sure the descriptor is known in the registry.
+        # Raises KeyError if unknown
+        pool.FindMessageTypeByName(type_name)
+
+        # create log
+        req_dict = {"@type": type_url, "policyDelta": {}}
+        req_struct = self._dict_to_struct(req_dict)
+
+        logger = Config.CLIENT.logger(f"auditdata-proto-{uuid.uuid1()}")
+        logger.log_proto(req_struct)
+
+        # retrieve log
+        retry = RetryErrors((TooManyRequests, StopIteration), max_tries=8)
+        protobuf_entry = retry(lambda: next(logger.list_entries()))()
+
+        self.assertIsInstance(protobuf_entry, entries.ProtobufEntry)
+        self.assertIsNone(protobuf_entry.payload_pb)
+        self.assertIsInstance(protobuf_entry.payload_json, dict)
+        self.assertEqual(protobuf_entry.payload_json["@type"], type_url)
+        self.assertEqual(
+            protobuf_entry.to_api_repr()["protoPayload"]["@type"], type_url
+        )
+
     def test_log_text(self):
         TEXT_PAYLOAD = "System test: test_log_text"
         logger = Config.CLIENT.logger(self._logger_name("log_text"))