From a0f97103b113fb0d954003882683ecdc47f08551 Mon Sep 17 00:00:00 2001 From: Guillaume Charest <1690085+gcharest@users.noreply.github.com> Date: Wed, 17 Apr 2024 17:24:14 +0000 Subject: [PATCH] feat: setup execute google api call function & tests --- .../google_workspace/google_service.py | 64 ++++++- .../google_workspace/test_google_service.py | 174 ++++++++++++++++++ 2 files changed, 233 insertions(+), 5 deletions(-) diff --git a/app/integrations/google_workspace/google_service.py b/app/integrations/google_workspace/google_service.py index c40a600e..48987552 100644 --- a/app/integrations/google_workspace/google_service.py +++ b/app/integrations/google_workspace/google_service.py @@ -1,7 +1,7 @@ """ Google Service Module. -This module provides a function to get an authenticated Google service and a decorator to handle Google API errors. +This module provides a function to get an authenticated Google service, a decorator to handle Google API errors, and a generic function to execute the Google API call. Functions: get_google_service(service: str, version: str) -> googleapiclient.discovery.Resource: @@ -23,9 +23,19 @@ from googleapiclient.errors import HttpError, Error # type: ignore from google.auth.exceptions import RefreshError # type: ignore +# Define the default arguments +DEFAULT_DELEGATED_ADMIN_EMAIL = os.environ.get("GOOGLE_DELEGATED_ADMIN_EMAIL") +DEFAULT_GOOGLE_WORKSPACE_CUSTOMER_ID = os.environ.get("GOOGLE_WORKSPACE_CUSTOMER_ID") + load_dotenv() +def convert_to_camel_case(snake_str): + """Convert a snake_case string to camelCase.""" + components = snake_str.split("_") + return components[0] + "".join(x.title() for x in components[1:]) + + def get_google_service(service, version, delegated_user_email=None, scopes=None): """ Get an authenticated Google service. @@ -76,15 +86,59 @@ def wrapper(*args, **kwargs): return func(*args, **kwargs) except HttpError as e: logging.error(f"An HTTP error occurred in function '{func.__name__}': {e}") - return None except ValueError as e: logging.error(f"A ValueError occurred in function '{func.__name__}': {e}") - return None except RefreshError as e: logging.error(f"A RefreshError occurred in function '{func.__name__}': {e}") - return None except Error as e: logging.error(f"An error occurred in function '{func.__name__}': {e}") - return None + except Exception as e: # Catch-all for any other types of exceptions + logging.error( + f"An unexpected error occurred in function '{func.__name__}': {e}" + ) + return None return wrapper + + +def execute_google_api_call( + service_name, + version, + resource, + method, + scopes=None, + delegated_user_email=None, + paginate=False, + **kwargs, +): + """Execute a Google API call. + + Args: + service_name (str): The name of the Google service. + version (str): The version of the Google service. + resource (str): The resource to access. + method (str): The method to call on the resource. + scopes (list, optional): The scopes for the Google service. + delegated_user_email (str, optional): The email address of the user to impersonate. + paginate (bool, optional): Whether to paginate the API call. + **kwargs: Additional keyword arguments for the API call. + + Returns: + dict or list: The result of the API call. If paginate is True, returns a list of all results. + """ + service = get_google_service(service_name, version, delegated_user_email, scopes) + resource_obj = getattr(service, resource)() + api_method = getattr(resource_obj, method) + if paginate: + all_results = [] + request = api_method(**kwargs) + while request is not None: + results = request.execute() + if results is not None: + all_results.extend( + results.get(resource, []) + ) # Use the resource name instead of "users" + request = getattr(resource_obj, method + "_next")(request, results) + return all_results + else: + return api_method(**kwargs).execute() diff --git a/app/tests/integrations/google_workspace/test_google_service.py b/app/tests/integrations/google_workspace/test_google_service.py index f1b80aca..0bb6768d 100644 --- a/app/tests/integrations/google_workspace/test_google_service.py +++ b/app/tests/integrations/google_workspace/test_google_service.py @@ -9,9 +9,19 @@ from integrations.google_workspace.google_service import ( get_google_service, handle_google_api_errors, + execute_google_api_call, + convert_to_camel_case, ) +def test_convert_to_camel_case(): + assert convert_to_camel_case("snake_case") == "snakeCase" + assert convert_to_camel_case("longer_snake_case_string") == "longerSnakeCaseString" + assert convert_to_camel_case("alreadyCamelCase") == "alreadyCamelCase" + assert convert_to_camel_case("singleword") == "singleword" + assert convert_to_camel_case("with_numbers_123") == "withNumbers123" + + @patch("integrations.google_workspace.google_service.build") @patch.object(Credentials, "from_service_account_info") def test_get_google_service_returns_build_object(credentials_mock, build_mock): @@ -29,6 +39,38 @@ def test_get_google_service_returns_build_object(credentials_mock, build_mock): ) +@patch("integrations.google_workspace.google_service.build") +@patch.object(Credentials, "from_service_account_info") +def test_get_google_service_with_delegated_user_email(credentials_mock, build_mock): + """ + Test case to verify that the function works correctly with a delegated user email. + """ + credentials_mock.return_value = MagicMock() + with patch.dict( + "os.environ", + {"GCP_SRE_SERVICE_ACCOUNT_KEY_FILE": json.dumps({"type": "service_account"})}, + ): + get_google_service("drive", "v3", delegated_user_email="test@test.com") + credentials_mock.return_value.with_subject.assert_called_once_with("test@test.com") + + +@patch("integrations.google_workspace.google_service.build") +@patch.object(Credentials, "from_service_account_info") +def test_get_google_service_with_scopes(credentials_mock, build_mock): + """ + Test case to verify that the function works correctly with scopes. + """ + credentials_mock.return_value = MagicMock() + with patch.dict( + "os.environ", + {"GCP_SRE_SERVICE_ACCOUNT_KEY_FILE": json.dumps({"type": "service_account"})}, + ): + get_google_service("drive", "v3", scopes=["scope1", "scope2"]) + credentials_mock.return_value.with_scopes.assert_called_once_with( + ["scope1", "scope2"] + ) + + def test_get_google_service_raises_exception_if_credentials_json_not_set(): """ Test case to verify that the function raises an exception if: @@ -102,6 +144,21 @@ def test_handle_google_api_errors_catches_error(mocked_logging_error): ) +@patch("logging.error") +def test_handle_google_api_errors_catches_exception(mocked_logging_error): + mock_func = MagicMock(side_effect=Exception("Exception message")) + mock_func.__name__ = "mock_func" + decorated_func = handle_google_api_errors(mock_func) + + result = decorated_func() + + assert result is None + mock_func.assert_called_once() + mocked_logging_error.assert_called_once_with( + "An unexpected error occurred in function 'mock_func': Exception message" + ) + + @patch("logging.error") def test_handle_google_api_errors_catches_refresh_error(mocked_logging_error): mock_func = MagicMock(side_effect=RefreshError("RefreshError message")) @@ -125,3 +182,120 @@ def test_handle_google_api_errors_passes_through_return_value(): assert result == "test" mock_func.assert_called_once() + + +@patch("integrations.google_workspace.google_service.get_google_service") +def test_execute_google_api_call_calls_get_google_service(mock_get_google_service): + execute_google_api_call("service_name", "version", "resource", "method") + mock_get_google_service.assert_called_once_with( + "service_name", "version", None, None + ) + + +@patch("integrations.google_workspace.google_service.get_google_service") +def test_execute_google_api_call_calls_get_google_service_with_delegated_user_email( + mock_get_google_service, +): + execute_google_api_call( + "service_name", + "version", + "resource", + "method", + delegated_user_email="admin.user@email.com", + ) + mock_get_google_service.assert_called_once_with( + "service_name", + "version", + "admin.user@email.com", + None, + ) + + +@patch("integrations.google_workspace.google_service.get_google_service") +def test_execute_google_api_call_calls_getattr_with_service_and_resource( + mock_get_google_service, +): + mock_service = MagicMock() + mock_get_google_service.return_value = mock_service + + execute_google_api_call("service_name", "version", "resource", "method") + + mock_service.resource.assert_called_once() + + +@patch("integrations.google_workspace.google_service.get_google_service") +def test_execute_google_api_call_when_paginate_is_false( + mock_get_google_service, +): + mock_service = MagicMock() + mock_get_google_service.return_value = mock_service + + # Set up the MagicMock for resource + mock_resource = MagicMock() + mock_service.resource.return_value = mock_resource + + mock_request = MagicMock() + mock_request.execute.return_value = {"key": "value"} + + # Set up the MagicMock for method + mock_resource.method.return_value = mock_request + + result = execute_google_api_call( + "service_name", "version", "resource", "method", arg1="value1" + ) + + mock_resource.method.assert_called_once_with(arg1="value1") + assert result == {"key": "value"} + + +@patch("integrations.google_workspace.google_service.get_google_service") +def test_execute_google_api_call_when_paginate_is_true( + mock_get_google_service, +): + mock_service = MagicMock() + mock_get_google_service.return_value = mock_service + + # Set up the MagicMock for resource + mock_resource = MagicMock() + mock_service.resource.return_value = mock_resource + + # Set up the MagicMock for request + mock_request1 = MagicMock() + mock_request1.execute.return_value = { + "resource": ["value1", "value2"], + "nextPageToken": "token", + } + + mock_request2 = MagicMock() + mock_request2.execute.return_value = {"resource": ["value3"], "nextPageToken": None} + + # Set up the MagicMock for method + mock_method = MagicMock() + mock_method.return_value = mock_request1 + mock_resource.method = mock_method + + # Set up the MagicMock for method_next + mock_method_next = MagicMock() + mock_resource.method_next = mock_method_next + + # Create a list of mock requests for pagination + mock_requests = [mock_request2] + + def side_effect(*args): + if mock_requests: + return mock_requests.pop(0) + else: + return None + + mock_method_next.side_effect = side_effect + + result = execute_google_api_call( + "service_name", "version", "resource", "method", paginate=True, arg1="value1" + ) + + assert result == ["value1", "value2", "value3"] + mock_resource.method.assert_called_once_with(arg1="value1") + mock_resource.method_next.assert_any_call( + mock_request1, {"resource": ["value1", "value2"], "nextPageToken": "token"} + ) + assert mock_method_next.call_count == 2