Retarus Secure Email Gateway

import uuid
from contextlib import contextmanager

import demistomock as demisto
import pytest
import RetarusSecureEmailGateway
from CommonServerPython import *
from RetarusSecureEmailGateway import (
    Connection,
    EventConnection,
    datetime,
    fetch_events,
    json,
    long_running_execution_command,
    timedelta,
)

CURRENT_TIME: datetime | None = None

EVENTS = [
    {"ts": "2023-08-16T13:24:12.147573+0100", "id": "1"},
    {"ts": "2023-08-14T13:24:12.147573+0200", "id": "2"},
    {"ts": "2023-08-12T13:24:11.147573+0000", "id": "3"},
]


def is_interval_passed(fetch_start_time: datetime, fetch_interval: int) -> bool:
    global CURRENT_TIME
    if not CURRENT_TIME:
        CURRENT_TIME = fetch_start_time
    return fetch_start_time + timedelta(seconds=fetch_interval) < CURRENT_TIME


@pytest.fixture
def connection():
    # Set up a mock connection
    return MockConnection()


class MockConnection(Connection):
    def __init__(
        self,
    ):
        global CURRENT_TIME
        self.id = uuid.uuid4()
        self.events = EVENTS
        self.index = 0
        self.pongs = 0
        self.create_time = datetime.now()

    def recv(self, timeout):
        global CURRENT_TIME
        # pretend to sleep for 4 seconds
        assert CURRENT_TIME
        CURRENT_TIME += timedelta(seconds=4)

        if self.index >= len(self.events):
            raise TimeoutError
        event = self.events[self.index]
        self.index += 1
        return json.dumps(event)

    def pong(self):
        self.pongs += 1


def test_heartbeat(mocker, connection):
    """
    Given:
        - A connection object with scarce messages

    When:
        - The long running execution loop runs

    Then:
        - Periodic keep-alive messages (pongs) are sent to the websocket connection to prevent it from closing.

    """
    idle_timeout = 3

    @contextmanager
    def mock_websocket_connection(host, cluster_id, api_key, since_time=None, to_time=None, fetch_interval=60):
        yield EventConnection(connection, fetch_interval, idle_timeout)

    def mock_perform_long_running_loop(connection, interval):
        # This mock will raise exceptions to stop the long running loop
        # StopIteration exception marks success
        if connection.connection.pongs:
            raise StopIteration(f"Sent {connection.connection.pongs} pongs")
        if datetime.now() > connection.create_time + timedelta(seconds=idle_timeout + 2):
            # Heartbeat should've been sent already
            raise TimeoutError(f"No heartbeat sent within {idle_timeout} seconds")

    mocker.patch.object(RetarusSecureEmailGateway, "websocket_connection", side_effect=mock_websocket_connection)
    mocker.patch.object(RetarusSecureEmailGateway, "perform_long_running_loop", side_effect=mock_perform_long_running_loop)

    with pytest.raises(StopIteration):
        long_running_execution_command(url="url", token_id="token_id", fetch_interval=60, channel="channel", verify_ssl=False)

    assert connection.pongs > 0


def test_fetch_events(mocker, connection):
    """
    Given:
        A connection to the websocket

    When:
        Calling fetch_events function to get events from the websocket connection

    Then:
        - Ensure that the function returns the events from the websocket connection
        - Ensure that the function converts the timestamp to UTC
        - Ensure that the function returns the events collected in the interval.
    """

    # We set fetch_interval to 7 to get this first two events (as we "wait" 4 seconds between each event)
    fetch_interval = 7
    event_connection = EventConnection(connection=connection)
    mocker.patch.object(RetarusSecureEmailGateway, "is_interval_passed", side_effect=is_interval_passed)
    debug_logs = mocker.patch.object(demisto, "debug")
    events = fetch_events(connection=event_connection, fetch_interval=fetch_interval)

    assert len(events) == 2
    assert events[0]["_time"] == "2023-08-16T12:24:12.147573+00:00"
    assert events[1]["_time"] == "2023-08-14T11:24:12.147573+00:00"

    debug_logs.assert_any_call("Retarus-logs Fetched 2 events")


def test_get_last_run_results_command__with_results(mocker):
    cnx = {"last_run_results": "results"}
    mocker.patch.object(demisto, "getIntegrationContext", return_value=cnx)
    res = RetarusSecureEmailGateway.get_last_run_results_command()
    assert res.readable_output == "results"


def test_get_last_run_results_command__no_results(mocker):
    mocker.patch.object(demisto, "getIntegrationContext", return_value={})
    res = RetarusSecureEmailGateway.get_last_run_results_command()
    assert "No results from the last run yet." in res.readable_output