fix: revert dict back to protobuf in the iam binding update (#1838)

* fix: update the iam binding update logic

This is needed as some recent updates to a dependency (google-cloud-resource-manager or its dependencies) broke the existing logic, and we are seeing the
error like this:
--> policy.bindings.append(new_binding)
TypeError: Expected a message object, but got {'role': 'roles/run.invoker', 'members': [...]}

* use right import

* add unit test mocking the iam update method

* use mock bq connection client

* import module rephrasing

Author: shobsi

bigframes/clients.py

@@ -24,6 +24,7 @@
 import google.api_core.exceptions
 import google.api_core.retry
 from google.cloud import bigquery_connection_v1, resourcemanager_v3
+from google.iam.v1 import policy_pb2
 
 logger = logging.getLogger(__name__)
 
@@ -172,10 +173,7 @@ def _ensure_iam_binding(
                     return
 
         # Create a new binding
-        new_binding = {
-            "role": role,
-            "members": [service_account],
-        }  # Use a dictionary to avoid problematic google.iam namespace package.
+        new_binding = policy_pb2.Binding(role=role, members=[service_account])
         policy.bindings.append(new_binding)
         request = {
             "resource": project,

tests/unit/test_clients.py

@@ -12,6 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest import mock
+
+from google.cloud import bigquery_connection_v1, resourcemanager_v3
+from google.iam.v1 import policy_pb2
 import pytest
 
 from bigframes import clients
@@ -65,3 +69,27 @@ def test_get_canonical_bq_connection_id_invalid_path():
             default_project="default-project",
             default_location="us",
         )
+
+
+def test_ensure_iam_binding():
+    bq_connection_client = mock.create_autospec(
+        bigquery_connection_v1.ConnectionServiceClient, instance=True
+    )
+    resource_manager_client = mock.create_autospec(
+        resourcemanager_v3.ProjectsClient, instance=True
+    )
+    resource_manager_client.get_iam_policy.return_value = policy_pb2.Policy(
+        bindings=[
+            policy_pb2.Binding(
+                role="roles/test.role1", members=["serviceAccount:serviceAccount1"]
+            )
+        ]
+    )
+    bq_connection_manager = clients.BqConnectionManager(
+        bq_connection_client, resource_manager_client
+    )
+    bq_connection_manager._IAM_WAIT_SECONDS = 0  # no need to wait in test
+    bq_connection_manager._ensure_iam_binding(
+        "test-project", "serviceAccount2", "roles/test.role2"
+    )
+    resource_manager_client.set_iam_policy.assert_called_once()