diff --git a/pulpcore/content/__init__.py b/pulpcore/content/__init__.py index 26d5200de4..f0c7720c5d 100644 --- a/pulpcore/content/__init__.py +++ b/pulpcore/content/__init__.py @@ -58,18 +58,18 @@ async def _heartbeat(): msg = "Content App '{name}' heartbeat written, sleeping for '{interarrival}' seconds".format( name=name, interarrival=heartbeat_interval ) - fail_msg = ("Content App '{name}' failed to write a heartbeat to the database.").format( - name=name - ) + fail_msg = f"Content App '{name}' failed to register or write a heartbeat to the database." versions = {app.label: app.version for app in pulp_plugin_configs()} - try: app_status = await AppStatus.objects.acreate( name=name, app_type="content", versions=versions ) except IntegrityError: log.error(f"A content app with name {name} already exists in the database.") - exit(Arbiter.WORKER_BOOT_ERROR) + os._exit(Arbiter.WORKER_BOOT_ERROR) + except (InterfaceError, DatabaseError) as e: + log.error(f"Content App '{name}' failed to register with the database. Exception: {str(e)}") + os._exit(Arbiter.WORKER_BOOT_ERROR) try: while True: await asyncio.sleep(heartbeat_interval) @@ -83,7 +83,7 @@ async def _heartbeat(): log.debug(msg) except (InterfaceError, DatabaseError) as e: log.error(f"{fail_msg} Exception: {str(e)}") - exit(Arbiter.WORKER_BOOT_ERROR) + os._exit(Arbiter.WORKER_BOOT_ERROR) finally: if app_status: await app_status.adelete() diff --git a/pulpcore/tests/unit/content/test_heartbeat.py b/pulpcore/tests/unit/content/test_heartbeat.py index 5fff207643..04fa6f1bf4 100644 --- a/pulpcore/tests/unit/content/test_heartbeat.py +++ b/pulpcore/tests/unit/content/test_heartbeat.py @@ -1,7 +1,8 @@ -from unittest.mock import AsyncMock, Mock, call +from unittest.mock import AsyncMock, Mock import pytest -from django.db.utils import InterfaceError, OperationalError +from django.db.utils import IntegrityError, InterfaceError, OperationalError +from gunicorn.arbiter import Arbiter from pulpcore.app.models.status import AppStatus, AppStatusManager from pulpcore.content import _heartbeat @@ -14,12 +15,7 @@ class MockException(Exception): @pytest.mark.parametrize("error_class", [InterfaceError, OperationalError]) @pytest.mark.asyncio -async def test_db_connection_interface_error(monkeypatch, settings, error_class): - """ - Test that if an InterfaceError or OperationalError is raised, - Handler._reset_db_connection() is called - """ - +async def test_exits_on_heartbeat_database_failure(monkeypatch, settings, error_class): mock_app_status = AsyncMock() mock_app_status.asave_heartbeat.side_effect = [error_class(), error_class()] mock_acreate = AsyncMock() @@ -28,10 +24,32 @@ async def test_db_connection_interface_error(monkeypatch, settings, error_class) monkeypatch.setattr(AppStatus, "objects", AppStatusManager()) mock_reset_db = Mock() monkeypatch.setattr(Handler, "_reset_db_connection", mock_reset_db) + mock_exit = Mock(side_effect=SystemExit(Arbiter.WORKER_BOOT_ERROR)) + monkeypatch.setattr("pulpcore.content.os._exit", mock_exit) settings.CONTENT_APP_TTL = 1 + monkeypatch.setattr("pulpcore.content.asyncio.sleep", AsyncMock(return_value=None)) + + with pytest.raises(SystemExit) as exc: + await _heartbeat() + + assert exc.value.code == Arbiter.WORKER_BOOT_ERROR + assert mock_app_status.asave_heartbeat.await_count == 2 + mock_reset_db.assert_called_once_with() + mock_exit.assert_called_once_with(Arbiter.WORKER_BOOT_ERROR) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("error_class", [IntegrityError, OperationalError]) +async def test_exits_on_initial_registration_database_failure(monkeypatch, settings, error_class): + mock_acreate = AsyncMock(side_effect=error_class()) + monkeypatch.setattr(AppStatusManager, "acreate", mock_acreate) + monkeypatch.setattr(AppStatus, "objects", AppStatusManager()) + mock_exit = Mock(side_effect=SystemExit(Arbiter.WORKER_BOOT_ERROR)) + monkeypatch.setattr("pulpcore.content.os._exit", mock_exit) - with pytest.raises(SystemExit): + with pytest.raises(SystemExit) as exc: await _heartbeat() - mock_app_status.asave_heartbeat.assert_called() - mock_reset_db.assert_has_calls([call()]) + assert exc.value.code == Arbiter.WORKER_BOOT_ERROR + mock_exit.assert_called_once_with(Arbiter.WORKER_BOOT_ERROR) + assert mock_acreate.await_count == 1