sientia-dataops-opc-ingestor/tests/unit/test_app.py

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 OsExitCalledError(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 OsExitCalledError(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'])

    mocks['GatewayConfig'] = MagicMock()
    mocks['GatewayRuntime'] = MagicMock()
    monkeypatch.setattr(app, 'GatewayConfig', mocks['GatewayConfig'])
    monkeypatch.setattr(app, 'GatewayRuntime', mocks['GatewayRuntime'])
    mock_runtime_instance = mocks['GatewayRuntime'].return_value
    mock_runtime_instance.start = AsyncMock()
    mock_runtime_instance.idle_sweep = AsyncMock()
    mock_runtime_instance.stop = AsyncMock()

    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()
    mock_ingestor_instance.ingestor_manager = 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(OsExitCalledError) 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.')

    # GatewayRuntime (Task 18): built once prepare_ingestor succeeds, started before
    # the loop, idle-swept each iteration, stopped on shutdown.
    mock_app_env['GatewayRuntime'].assert_called_once_with(
        config=mock_app_env['GatewayConfig'].from_env.return_value,
        data_manager=mock_ingestor_instance.ingestor_manager.data_manager,
        pool=mock_ingestor_instance.ingestor_manager.pool,
        logger=mock_ingestor_instance.logger,
        notification_handler=mock_ingestor_instance.notification_handler,
        metrics_controller=mock_ingestor_instance.metrics_controller,
        metadata=mock_ingestor_instance.metadata,
    )
    mock_runtime_instance = mock_app_env['GatewayRuntime'].return_value
    mock_runtime_instance.start.assert_awaited_once()
    mock_runtime_instance.idle_sweep.assert_awaited_once()
    mock_runtime_instance.stop.assert_awaited_once()

    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(OsExitCalledError) 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(OsExitCalledError) 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(OsExitCalledError) 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(OsExitCalledError) 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(OsExitCalledError):
        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(OsExitCalledError) 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_ingestor_instance.ingestor_manager = None  # prepare failed before assignment
    mock_exit_signal.is_set.side_effect = [False, True]

    with pytest.raises(OsExitCalledError) as excinfo:
        await app.main()
    assert excinfo.value.code == 0

    # GatewayRuntime must not be built when prepare_ingestor never got a data_manager
    mock_app_env['GatewayRuntime'].assert_not_called()

    # 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'
    )


@pytest.mark.asyncio
async def test_main_gateway_runtime_init_failure_keeps_read_path_alive(mock_app_env):
    """GatewayRuntime construction failing (e.g. Postgres config) must not kill the
    app: the read path lives in the ingestor_manager, so the main loop still runs and
    shuts down cleanly."""
    mock_ingestor_instance = mock_app_env['Ingestor'].return_value
    mock_exit_signal = mock_app_env['mock_exit_signal']

    mock_app_env['GatewayRuntime'].side_effect = Exception('MinIO unreachable')
    mock_exit_signal.is_set.side_effect = [False, True]
    mock_app_env['time_time'].side_effect = [10.0, 11.0]

    with pytest.raises(OsExitCalledError) as excinfo:
        await app.main()
    assert excinfo.value.code == 0

    # Loop still ran and shutdown was clean, without any runtime calls
    mock_ingestor_instance.loop.assert_called_once()
    mock_ingestor_instance.shutdown.assert_called_once()
    mock_runtime_instance = mock_app_env['GatewayRuntime'].return_value
    mock_runtime_instance.idle_sweep.assert_not_called()
    mock_runtime_instance.stop.assert_not_called()

    # Failure was logged and counted, but did not set exit_signal
    assert any(
        'GatewayRuntime failed to initialize' in str(c)
        for c in mock_ingestor_instance.logger.error.call_args_list
    )
    mock_app_env['metrics_APP_ERRORS_TOTAL_labels_inc'].assert_called_once()
    mock_exit_signal.set.assert_not_called()