Fix health check query not executed

CHANGELOG.md

@@ -12,6 +12,7 @@ ENHANCEMENTS:
 * Update API, CLI, and UI dependencies to address high-severity Dependabot alerts, including `PyJWT`, `Vite`, `lodash`, `fast-uri`, `flatted`, `immutable`, and `minimatch`.
 
 BUG FIXES:
+* Fix Health check can falsely return OK even if Cosmos is down or inaccessible. ([#4926](https://github.com/microsoft/AzureTRE/issues/4926))
 * Fix API timeout and name collision failures on workspace creation by checking storage account name availability and improved logging. ([#4946](https://github.com/microsoft/AzureTRE/pull/4946))
 * Fix error handling in airlock processor ([#4929](https://github.com/microsoft/AzureTRE/pull/4929))
 * Fix dependabot high severity alerts for packages fast-uri, lodash, picomatch, immutable, minimatch, flatted and PyJWT

api_app/_version.py

@@ -1 +1 @@
-__version__ = "0.25.26"
+__version__ = "0.25.27"

api_app/services/health_checker.py

@@ -19,7 +19,9 @@ async def create_state_store_status() -> Tuple[StatusEnum, str]:
     message = ""
     try:
         container: ContainerProxy = await Database().get_container_proxy(STATE_STORE_RESOURCES_CONTAINER)
-        container.query_items("SELECT TOP 1 * FROM c")
+        await container.read()
+        async for _ in container.query_items("SELECT TOP 1 * FROM c", max_item_count=1):
+            break
     except exceptions.ServiceRequestError:
         status = StatusEnum.not_ok
         message = strings.STATE_STORE_ENDPOINT_NOT_RESPONDING

api_app/tests_ma/test_services/test_health_checker.py

@@ -2,6 +2,7 @@
 from unittest.mock import AsyncMock, MagicMock
 import pytest
 from azure.core.exceptions import ServiceRequestError
+from azure.cosmos.exceptions import CosmosHttpResponseError
 from azure.servicebus.exceptions import ServiceBusConnectionError
 from mock import patch
 from models.schemas.status import StatusEnum
@@ -11,12 +12,48 @@
 pytestmark = pytest.mark.asyncio
 
 
-@patch("azure.cosmos.aio.ContainerProxy.query_items", return_value=AsyncMock())
-async def test_get_state_store_status_responding(_) -> None:
+def create_mock_container(query_results=None, query_error=None, read_error=None):
+    container_mock = MagicMock()
+
+    # Mock read()
+    if read_error:
+        container_mock.read = AsyncMock(side_effect=read_error)
+    else:
+        container_mock.read = AsyncMock(return_value={"id": "container_properties"})
+
+    # Mock query_items()
+    query_items_mock = MagicMock()
+    if query_error:
+        query_items_mock.return_value.__aiter__.side_effect = query_error
+    else:
+        query_items_mock.return_value.__aiter__.return_value = query_results or []
+
+    container_mock.query_items = query_items_mock
+    return container_mock
+
+
+@patch("api.dependencies.database.Database.get_container_proxy")
+async def test_get_state_store_status_responding(get_container_proxy_mock) -> None:
+    container_mock = create_mock_container(query_results=[{"id": "item"}])
+    get_container_proxy_mock.return_value = container_mock
+    status, message = await health_checker.create_state_store_status()
+
+    assert status == StatusEnum.ok
+    assert message == ""
+    container_mock.read.assert_called_once()
+    container_mock.query_items.assert_called_once_with("SELECT TOP 1 * FROM c", max_item_count=1)
+
+
+@patch("api.dependencies.database.Database.get_container_proxy")
+async def test_get_state_store_status_empty_results(get_container_proxy_mock) -> None:
+    container_mock = create_mock_container(query_results=[])
+    get_container_proxy_mock.return_value = container_mock
     status, message = await health_checker.create_state_store_status()
 
     assert status == StatusEnum.ok
     assert message == ""
+    container_mock.read.assert_called_once()
+    container_mock.query_items.assert_called_once_with("SELECT TOP 1 * FROM c", max_item_count=1)
 
 
 @patch("api.dependencies.database.Database.get_container_proxy")
@@ -39,6 +76,54 @@ async def test_get_state_store_status_other_exception(container_proxy_mock) -> N
     assert message == strings.UNSPECIFIED_ERROR
 
 
+@patch("api.dependencies.database.Database.get_container_proxy")
+async def test_get_state_store_status_read_cosmos_http_error(get_container_proxy_mock) -> None:
+    container_mock = create_mock_container(read_error=CosmosHttpResponseError(message="some message"))
+    get_container_proxy_mock.return_value = container_mock
+    status, message = await health_checker.create_state_store_status()
+
+    assert status == StatusEnum.not_ok
+    assert message == strings.STATE_STORE_ENDPOINT_NOT_ACCESSIBLE
+    container_mock.read.assert_called_once()
+    container_mock.query_items.assert_not_called()
+
+
+@patch("api.dependencies.database.Database.get_container_proxy")
+async def test_get_state_store_status_read_service_request_error(get_container_proxy_mock) -> None:
+    container_mock = create_mock_container(read_error=ServiceRequestError(message="some message"))
+    get_container_proxy_mock.return_value = container_mock
+    status, message = await health_checker.create_state_store_status()
+
+    assert status == StatusEnum.not_ok
+    assert message == strings.STATE_STORE_ENDPOINT_NOT_RESPONDING
+    container_mock.read.assert_called_once()
+    container_mock.query_items.assert_not_called()
+
+
+@patch("api.dependencies.database.Database.get_container_proxy")
+async def test_get_state_store_status_cosmos_http_error(get_container_proxy_mock) -> None:
+    container_mock = create_mock_container(query_error=CosmosHttpResponseError(message="some message"))
+    get_container_proxy_mock.return_value = container_mock
+    status, message = await health_checker.create_state_store_status()
+
+    assert status == StatusEnum.not_ok
+    assert message == strings.STATE_STORE_ENDPOINT_NOT_ACCESSIBLE
+    container_mock.read.assert_called_once()
+    container_mock.query_items.assert_called_once_with("SELECT TOP 1 * FROM c", max_item_count=1)
+
+
+@patch("api.dependencies.database.Database.get_container_proxy")
+async def test_get_state_store_status_service_request_error(get_container_proxy_mock) -> None:
+    container_mock = create_mock_container(query_error=ServiceRequestError(message="some message"))
+    get_container_proxy_mock.return_value = container_mock
+    status, message = await health_checker.create_state_store_status()
+
+    assert status == StatusEnum.not_ok
+    assert message == strings.STATE_STORE_ENDPOINT_NOT_RESPONDING
+    container_mock.read.assert_called_once()
+    container_mock.query_items.assert_called_once_with("SELECT TOP 1 * FROM c", max_item_count=1)
+
+
 @patch("core.credentials.get_credential_async_context")
 @patch("services.health_checker.ServiceBusClient")
 async def test_get_service_bus_status_responding(service_bus_client_mock, get_credential_async_context) -> None: