import asyncio import signal as signal_module # To avoid conflict with mock names from threading import Event from unittest.mock import AsyncMock, MagicMock, call import pytest # Import the 'app' module to be tested from ingestor import app # Custom exception to catch os._exit calls class OsExitCalled(Exception): def __init__(self, code): super().__init__(f'os._exit({code}) called') self.code = code # Helper function for the os_exit mock's side_effect def raise_os_exit_with_code(exit_code): raise OsExitCalled(exit_code) @pytest.fixture def mock_app_env(monkeypatch): """Fixture to mock dependencies of app.main and app.signal_handler.""" mocks = { 'start_http_server': MagicMock(), 'Ingestor': MagicMock(), 'metrics_APP_UP_labels_set': MagicMock(), 'metrics_APP_LOOP_COUNT_labels_inc': MagicMock(), 'metrics_APP_LOOP_DURATION_labels_observe': MagicMock(), 'metrics_APP_ERRORS_TOTAL_labels_inc': MagicMock(), 'os_exit': MagicMock(side_effect=raise_os_exit_with_code), 'time_time': MagicMock(), 'asyncio_sleep': AsyncMock(), 'signal_signal': MagicMock(), 'traceback_print_exc': MagicMock(), 'mock_exit_signal': MagicMock(spec=Event), } monkeypatch.setattr(app, 'start_http_server', mocks['start_http_server']) monkeypatch.setattr(app, 'Ingestor', mocks['Ingestor']) monkeypatch.setattr(asyncio, 'sleep', mocks['asyncio_sleep']) monkeypatch.setattr( app.metrics.APP_UP, 'labels', MagicMock(return_value=MagicMock(set=mocks['metrics_APP_UP_labels_set'])), ) monkeypatch.setattr( app.metrics.APP_LOOP_COUNT, 'labels', MagicMock(return_value=MagicMock(inc=mocks['metrics_APP_LOOP_COUNT_labels_inc'])), ) monkeypatch.setattr( app.metrics.APP_LOOP_DURATION, 'labels', MagicMock( return_value=MagicMock(observe=mocks['metrics_APP_LOOP_DURATION_labels_observe']) ), ) monkeypatch.setattr( app.metrics.APP_ERRORS_TOTAL, 'labels', MagicMock(return_value=MagicMock(inc=mocks['metrics_APP_ERRORS_TOTAL_labels_inc'])), ) monkeypatch.setattr(app.os, '_exit', mocks['os_exit']) monkeypatch.setattr(app, 'time', mocks['time_time']) monkeypatch.setattr(app.signal, 'signal', mocks['signal_signal']) monkeypatch.setattr(app.traceback, 'print_exc', mocks['traceback_print_exc']) monkeypatch.setattr(app, 'exit_signal', mocks['mock_exit_signal']) monkeypatch.setattr(app, 'POD_ID', 'test_pod') mock_ingestor_instance = mocks['Ingestor'].return_value mock_ingestor_instance.poll_interval = 0.01 mock_ingestor_instance.logger = MagicMock() # Make async methods async mocks mock_ingestor_instance.prepare_ingestor = AsyncMock() mock_ingestor_instance.loop = AsyncMock() mock_ingestor_instance.shutdown = AsyncMock() return mocks @pytest.mark.asyncio async def test_main_successful_run_one_loop(mock_app_env, capsys): """Test a successful run where the loop executes once and then exits gracefully.""" mock_ingestor_instance = mock_app_env['Ingestor'].return_value mock_exit_signal = mock_app_env['mock_exit_signal'] mock_exit_signal.is_set.side_effect = [False, True] mock_app_env['time_time'].side_effect = [10.0, 11.5] with pytest.raises(OsExitCalled) as excinfo: await app.main() assert excinfo.value.code == 0 mock_app_env['start_http_server'].assert_called_once_with(9090) app.metrics.APP_UP.labels.assert_any_call(pod_id='test_pod') set_calls = mock_app_env['metrics_APP_UP_labels_set'].call_args_list assert call(1) in set_calls assert call(0) in set_calls assert set_calls.index(call(1)) < set_calls.index(call(0)) mock_app_env['Ingestor'].assert_called_once_with() mock_ingestor_instance.prepare_ingestor.assert_called_once() mock_ingestor_instance.logger.info.assert_any_call('Ingestor prepared. Starting main loop.') mock_ingestor_instance.loop.assert_called_once() app.metrics.APP_LOOP_COUNT.labels.assert_called_with(pod_id='test_pod') mock_app_env['metrics_APP_LOOP_COUNT_labels_inc'].assert_called_once() mock_app_env['asyncio_sleep'].assert_has_calls( [call(mock_ingestor_instance.poll_interval), call(5)] ) app.metrics.APP_LOOP_DURATION.labels.assert_called_with(pod_id='test_pod') mock_app_env['metrics_APP_LOOP_DURATION_labels_observe'].assert_called_once_with(1.5) mock_ingestor_instance.shutdown.assert_called_once() mock_ingestor_instance.logger.info.assert_any_call('Main loop exit_signaled.') captured = capsys.readouterr() assert 'Prometheus server started on port 9090.' in captured.out @pytest.mark.asyncio async def test_main_prometheus_server_fails_to_start(mock_app_env, capsys): """Test the scenario where starting the Prometheus server fails.""" mock_app_env['start_http_server'].side_effect = OSError('Port already in use') with pytest.raises(OsExitCalled) as excinfo: await app.main() assert excinfo.value.code == 1 # Check that .set(1) was not called. .set(0) definitely not called. called_with_1 = False for call_args in mock_app_env['metrics_APP_UP_labels_set'].call_args_list: if call_args == call(1): called_with_1 = True break assert not called_with_1, 'APP_UP.set(1) should not have been called if server start failed' mock_app_env['Ingestor'].assert_not_called() captured = capsys.readouterr() assert 'Failed to start Prometheus server: Port already in use' in captured.out @pytest.mark.asyncio async def test_main_loop_exception_handling(mock_app_env, capsys): """Test that an exception in ingestor.loop() is handled gracefully.""" mock_ingestor_instance = mock_app_env['Ingestor'].return_value mock_exit_signal = mock_app_env['mock_exit_signal'] mock_exit_signal.is_set.side_effect = [False, True] mock_ingestor_instance.loop.side_effect = Exception('Test loop exception') mock_app_env['time_time'].side_effect = [10.0, 10.1] with pytest.raises(OsExitCalled) as excinfo: await app.main() assert excinfo.value.code == 0 mock_ingestor_instance.loop.assert_called_once() mock_app_env['traceback_print_exc'].assert_called_once() app.metrics.APP_ERRORS_TOTAL.labels.assert_called_with(pod_id='test_pod') mock_app_env['metrics_APP_ERRORS_TOTAL_labels_inc'].assert_called_once() mock_exit_signal.set.assert_called_once() mock_app_env['metrics_APP_LOOP_COUNT_labels_inc'].assert_not_called() app.metrics.APP_LOOP_DURATION.labels.assert_called_with(pod_id='test_pod') mock_app_env['metrics_APP_LOOP_DURATION_labels_observe'].assert_called_once_with( pytest.approx(0.1) ) mock_ingestor_instance.shutdown.assert_called_once() captured = capsys.readouterr() assert 'Exception in main loop. Setting exit_signal flag.' in captured.out @pytest.mark.asyncio async def test_main_keyboard_interrupt_handling(mock_app_env, capsys): """Test that KeyboardInterrupt in ingestor.loop() is handled.""" mock_ingestor_instance = mock_app_env['Ingestor'].return_value mock_exit_signal = mock_app_env['mock_exit_signal'] mock_exit_signal.is_set.side_effect = [False, True] mock_ingestor_instance.loop.side_effect = KeyboardInterrupt() mock_app_env['time_time'].side_effect = [10.0, 10.1] with pytest.raises(OsExitCalled) as excinfo: await app.main() assert excinfo.value.code == 0 mock_ingestor_instance.loop.assert_called_once() mock_exit_signal.set.assert_called_once() mock_ingestor_instance.shutdown.assert_called_once() captured = capsys.readouterr() assert 'KeyboardInterrupt received. Setting exit_signal flag.' in captured.out mock_app_env['metrics_APP_ERRORS_TOTAL_labels_inc'].assert_not_called() def test_signal_handler_sets_exit_signal(mock_app_env): """Test that the signal_handler function calls exit_signal.set().""" mock_exit_signal_set = mock_app_env['mock_exit_signal'].set app.signal_handler(signal_module.SIGINT, None) mock_exit_signal_set.assert_called_once() @pytest.mark.asyncio async def test_main_multiple_loop_iterations(mock_app_env): """Test the main loop runs for a few iterations.""" mock_ingestor_instance = mock_app_env['Ingestor'].return_value mock_exit_signal = mock_app_env['mock_exit_signal'] mock_exit_signal.is_set.side_effect = [False, False, False, True] mock_app_env['time_time'].side_effect = [10.0, 10.1, 10.2, 10.3, 10.4, 10.5] with pytest.raises(OsExitCalled) as excinfo: await app.main() assert excinfo.value.code == 0 assert mock_ingestor_instance.loop.call_count == 3 assert app.metrics.APP_LOOP_COUNT.labels.call_count == 3 app.metrics.APP_LOOP_COUNT.labels.assert_called_with( pod_id='test_pod' ) # Checks last call or any call assert mock_app_env['metrics_APP_LOOP_COUNT_labels_inc'].call_count == 3 assert mock_app_env['asyncio_sleep'].call_count == 4 assert app.metrics.APP_LOOP_DURATION.labels.call_count == 3 app.metrics.APP_LOOP_DURATION.labels.assert_called_with(pod_id='test_pod') duration_calls = mock_app_env['metrics_APP_LOOP_DURATION_labels_observe'].call_args_list assert duration_calls[0] == call(pytest.approx(0.1, abs=1e-9)) assert duration_calls[1] == call(pytest.approx(0.1, abs=1e-9)) assert duration_calls[2] == call(pytest.approx(0.1, abs=1e-9)) mock_ingestor_instance.shutdown.assert_called_once() @pytest.mark.asyncio async def test_main_pod_id_used_in_metrics(mock_app_env): """Test that the POD_ID from app module is used in metric labels.""" mock_exit_signal = mock_app_env['mock_exit_signal'] mock_exit_signal.is_set.side_effect = [False, True] mock_app_env['time_time'].side_effect = [10.0, 11.0] with pytest.raises(OsExitCalled): await app.main() app.metrics.APP_UP.labels.assert_any_call(pod_id='test_pod') app.metrics.APP_LOOP_COUNT.labels.assert_any_call(pod_id='test_pod') app.metrics.APP_LOOP_DURATION.labels.assert_any_call(pod_id='test_pod') # APP_ERRORS_TOTAL would be checked similarly if it were called in this flow. # Check the .set() / .inc() calls on the mocks returned by .labels() mock_app_env['metrics_APP_UP_labels_set'].assert_any_call(1) mock_app_env['metrics_APP_UP_labels_set'].assert_any_call(0) mock_app_env['metrics_APP_LOOP_COUNT_labels_inc'].assert_called_once() @pytest.mark.asyncio async def test_run_async_main(mock_app_env, capsys): """Test the run_async_main function that sets up the event loop.""" mock_ingestor_instance = mock_app_env['Ingestor'].return_value mock_exit_signal = mock_app_env['mock_exit_signal'] mock_exit_signal.is_set.side_effect = [False, True] mock_app_env['time_time'].side_effect = [10.0, 11.0] # Instead of calling run_async_main() which creates a new event loop, # we test the main() function directly since that's what run_async_main() would call with pytest.raises(OsExitCalled) as excinfo: await app.main() assert excinfo.value.code == 0 # Verify the main function was called through the event loop mock_app_env['start_http_server'].assert_called_once_with(9090) mock_ingestor_instance.prepare_ingestor.assert_called_once() mock_ingestor_instance.loop.assert_called_once() mock_ingestor_instance.shutdown.assert_called_once() @pytest.mark.asyncio async def test_main_prepare_ingestor_failure(mock_app_env, capsys): """Test that prepare_ingestor failure is handled correctly.""" mock_ingestor_instance = mock_app_env['Ingestor'].return_value mock_exit_signal = mock_app_env['mock_exit_signal'] mock_ingestor_instance.prepare_ingestor.side_effect = Exception('Preparation failed') mock_exit_signal.is_set.side_effect = [False, True] with pytest.raises(OsExitCalled) as excinfo: await app.main() assert excinfo.value.code == 0 # Verify error metrics were incremented app.metrics.APP_ERRORS_TOTAL.labels.assert_called_with(pod_id='test_pod') mock_app_env['metrics_APP_ERRORS_TOTAL_labels_inc'].assert_called_once() # Verify exit signal was set mock_exit_signal.set.assert_called_once() # Verify shutdown was called mock_ingestor_instance.shutdown.assert_called_once() # Verify error was logged mock_ingestor_instance.logger.error.assert_called_once_with( 'Failed to prepare ingestor: Preparation failed' )