diff --git a/app/integrations/aws/client.py b/app/integrations/aws/client.py index 4a510205..66ba5b77 100644 --- a/app/integrations/aws/client.py +++ b/app/integrations/aws/client.py @@ -135,12 +135,3 @@ def paginator(client, operation, keys=None, **kwargs): results.extend(page[key]) return results - - -def healthcheck(): - """Check the health of the AWS integration. - - Returns: - bool: True if the integration is healthy, False otherwise. - """ - return execute_aws_api_call("sts", "get_caller_identity") is not False diff --git a/app/integrations/aws/identity_store.py b/app/integrations/aws/identity_store.py index fbf8b165..f403edf2 100644 --- a/app/integrations/aws/identity_store.py +++ b/app/integrations/aws/identity_store.py @@ -24,6 +24,22 @@ def resolve_identity_store_id(kwargs): return kwargs +def healthcheck(): + """Check the health of the AWS integration. + + Returns: + bool: True if the integration is healthy, False otherwise. + """ + healthy = False + try: + response = list_users() + healthy = True if response else False + logger.info(f"AWS IdentityStore healthcheck result: {response}") + except Exception as error: + logger.error(f"AWS IdentityStore healthcheck failed: {error}") + return healthy + + @handle_aws_api_errors def create_user(email, first_name, family_name, **kwargs): """Creates a new user in the AWS Identity Center (identitystore) diff --git a/app/jobs/scheduled_tasks.py b/app/jobs/scheduled_tasks.py index ea62de5c..1b8f0753 100644 --- a/app/jobs/scheduled_tasks.py +++ b/app/jobs/scheduled_tasks.py @@ -7,7 +7,7 @@ from integrations import google_drive, maxmind, opsgenie -from integrations.aws import client as aws_client +from integrations.aws import identity_store from modules.aws import identity_center logging.basicConfig(level=logging.INFO) @@ -36,7 +36,7 @@ def integration_healthchecks(): "google_drive": google_drive.healthcheck, "maxmind": maxmind.healthcheck, "opsgenie": opsgenie.healthcheck, - "aws": aws_client.healthcheck, + "aws": identity_store.healthcheck, } for key, healthcheck in healthchecks.items(): if not healthcheck(): diff --git a/app/tests/integrations/aws/test_client.py b/app/tests/integrations/aws/test_client.py index 07ebce83..5876d76b 100644 --- a/app/tests/integrations/aws/test_client.py +++ b/app/tests/integrations/aws/test_client.py @@ -312,23 +312,3 @@ def test_execute_aws_api_call_raises_exception_when_role_arn_not_provided( ) mock_assume_role.assert_not_called() mock_convert_kwargs_to_pascal_case.assert_not_called() - - -@patch("integrations.aws.client.execute_aws_api_call") -def test_healtcheck_is_healthy(mock_execute_aws_api_call): - mock_execute_aws_api_call.return_value = {"key": "value"} - - result = aws_client.healthcheck() - - assert result is True - mock_execute_aws_api_call.assert_called_once_with("sts", "get_caller_identity") - - -@patch("integrations.aws.client.execute_aws_api_call") -def test_healtcheck_is_unhealthy(mock_execute_aws_api_call): - mock_execute_aws_api_call.return_value = False - - result = aws_client.healthcheck() - - assert result is False - mock_execute_aws_api_call.assert_called_once_with("sts", "get_caller_identity") diff --git a/app/tests/integrations/aws/test_identity_store.py b/app/tests/integrations/aws/test_identity_store.py index 119a18b1..27a6c669 100644 --- a/app/tests/integrations/aws/test_identity_store.py +++ b/app/tests/integrations/aws/test_identity_store.py @@ -55,6 +55,26 @@ def test_resolve_identity_store_id_no_env(): identity_store.resolve_identity_store_id({}) +@patch("integrations.aws.identity_store.list_users") +def test_healtcheck_is_healthy(mock_list_users): + mock_list_users.return_value = ["User1", "User2"] + + result = identity_store.healthcheck() + + assert result is True + mock_list_users.assert_called_once + + +@patch("integrations.aws.identity_store.list_users") +def test_healtcheck_is_unhealthy(mock_list_users): + mock_list_users.return_value = [] + + result = identity_store.healthcheck() + + assert result is False + mock_list_users.assert_called_once + + @patch("integrations.aws.identity_store.execute_aws_api_call") @patch("integrations.aws.identity_store.resolve_identity_store_id") def test_create_user(mock_resolve_identity_store_id, mock_execute_aws_api_call): diff --git a/app/tests/jobs/test_scheduled_tasks.py b/app/tests/jobs/test_scheduled_tasks.py index 7be0db35..ec6c5fe1 100644 --- a/app/tests/jobs/test_scheduled_tasks.py +++ b/app/tests/jobs/test_scheduled_tasks.py @@ -24,7 +24,7 @@ def test_run_continuously(time_mock, threading_mock, schedule_mock): assert result == cease_continuous_run -@patch("jobs.scheduled_tasks.aws_client") +@patch("jobs.scheduled_tasks.identity_store") @patch("jobs.scheduled_tasks.google_drive") @patch("jobs.scheduled_tasks.maxmind") @patch("jobs.scheduled_tasks.opsgenie") @@ -43,7 +43,7 @@ def test_integration_healthchecks_healthy( assert mock_logging.error.call_count == 0 -@patch("jobs.scheduled_tasks.aws_client") +@patch("jobs.scheduled_tasks.identity_store") @patch("jobs.scheduled_tasks.google_drive") @patch("jobs.scheduled_tasks.maxmind") @patch("jobs.scheduled_tasks.opsgenie")