fix: use list users as healthcheck (#510)

* fix: use list users as healthcheck

* fix: lint

Author: gcharest

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

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)

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():

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")

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):

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")