|
12 | 12 | # See the License for the specific language governing permissions and
|
13 | 13 | # limitations under the License.
|
14 | 14 |
|
| 15 | +from unittest import mock |
| 16 | + |
| 17 | +from google.cloud import bigquery_connection_v1, resourcemanager_v3 |
| 18 | +from google.iam.v1 import policy_pb2 |
15 | 19 | import pytest
|
16 | 20 |
|
17 | 21 | from bigframes import clients
|
@@ -65,3 +69,27 @@ def test_get_canonical_bq_connection_id_invalid_path():
|
65 | 69 | default_project="default-project",
|
66 | 70 | default_location="us",
|
67 | 71 | )
|
| 72 | + |
| 73 | + |
| 74 | +def test_ensure_iam_binding(): |
| 75 | + bq_connection_client = mock.create_autospec( |
| 76 | + bigquery_connection_v1.ConnectionServiceClient, instance=True |
| 77 | + ) |
| 78 | + resource_manager_client = mock.create_autospec( |
| 79 | + resourcemanager_v3.ProjectsClient, instance=True |
| 80 | + ) |
| 81 | + resource_manager_client.get_iam_policy.return_value = policy_pb2.Policy( |
| 82 | + bindings=[ |
| 83 | + policy_pb2.Binding( |
| 84 | + role="roles/test.role1", members=["serviceAccount:serviceAccount1"] |
| 85 | + ) |
| 86 | + ] |
| 87 | + ) |
| 88 | + bq_connection_manager = clients.BqConnectionManager( |
| 89 | + bq_connection_client, resource_manager_client |
| 90 | + ) |
| 91 | + bq_connection_manager._IAM_WAIT_SECONDS = 0 # no need to wait in test |
| 92 | + bq_connection_manager._ensure_iam_binding( |
| 93 | + "test-project", "serviceAccount2", "roles/test.role2" |
| 94 | + ) |
| 95 | + resource_manager_client.set_iam_policy.assert_called_once() |
0 commit comments