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

import pytest
from unittest.mock import patch, MagicMock, call
from threading import Event
import signal as signal_module  # To avoid conflict with mock names
import os

# 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),  # CORRECTED
        "time_time": MagicMock(),
        "time_sleep": MagicMock(),
        "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(
        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, "sleep", mocks["time_sleep"])
    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()

    return mocks


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:
        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_exit_signal.wait.assert_called_once_with(
        mock_ingestor_instance.poll_interval)

    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_app_env["time_sleep"].assert_called_once_with(5)
    # CORREÇÃO APLICADA ABAIXO:
    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


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:
        app.main()
    assert excinfo.value.code == 1

    # Verify that APP_UP.labels(...).set(1) was NOT called.
    # The mock for .set is mock_app_env['metrics_APP_UP_labels_set']
    # We need to check if it was called with 1.
    # A more robust way is to check if the specific .labels(pod_id="test_pod") mock was ever called
    # and then its .set(1) method.
    # For simplicity here, we check if metrics_APP_UP_labels_set was ever called with 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


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


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


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:
        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_exit_signal.wait.call_count == 3

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


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