daena/tests/test_plugins.py

from __future__ import annotations

import asyncio
import json
from pathlib import Path
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch

import pytest

from daena.pipeline.processors.enricher import EnricherProcessor
from daena.pipeline.processors.filter import FilterProcessor
from daena.pipeline.record import Record
from daena.pipeline.sinks.loki import LokiSink
from daena.pipeline.sources.journald import (
    JournaldSource,
    build_journalctl_command,
    entry_to_record,
)


def test_journalctl_command_tails_new_entries_without_cursor() -> None:
    command = build_journalctl_command(cursor=None, units=["ssh.service"], matches=["_PID=1"])
    assert command[:4] == ["journalctl", "--follow", "--output=json", "--no-pager"]
    assert "--lines=0" in command
    assert command[command.index("--unit") + 1] == "ssh.service"
    assert "_PID=1" in command


def test_journalctl_command_uses_cursor() -> None:
    command = build_journalctl_command(cursor="abc", units=[], matches=[])
    assert "--after-cursor=abc" in command
    assert "--lines=0" not in command


def test_entry_to_record_uses_journal_timestamp() -> None:
    record = entry_to_record(
        "system",
        {
            "MESSAGE": "hello",
            "SYSLOG_IDENTIFIER": "sshd",
            "_SYSTEMD_UNIT": "ssh.service",
            "__CURSOR": "cur",
            "__REALTIME_TIMESTAMP": "1600000000000000",
        },
    )
    assert record.source == "system"
    assert record.record_type == "journald"
    assert record.body["message"] == "hello"
    assert record.body["syslog_identifier"] == "sshd"
    assert record.metadata["unit"] == "ssh.service"
    assert record.timestamp.startswith("2020-")


class _ListReader:
    def __init__(self, entries: list[dict[str, Any]]) -> None:
        self._entries = entries
        self.stopped = False

    async def run(self, on_entry) -> None:  # type: ignore[no-untyped-def]
        for entry in self._entries:
            await on_entry(entry)

    async def stop(self) -> None:
        self.stopped = True


@pytest.mark.asyncio
async def test_journald_source_emits_and_stores_cursor(tmp_path: Path) -> None:
    cursor_path = tmp_path / "cursor"
    reader = _ListReader(
        [
            {"MESSAGE": "one", "__CURSOR": "c1"},
            {"MESSAGE": "two", "__CURSOR": "c2"},
        ]
    )
    source = JournaldSource(
        "journal",
        config={"cursor_path": str(cursor_path)},
        reader=reader,
    )
    seen: list[Record] = []

    async def capture(record: Record) -> None:
        seen.append(record)

    source.set_callback(capture)
    await source.start()
    await asyncio.sleep(0.05)
    await source.stop()
    assert [record.body["message"] for record in seen] == ["one", "two"]
    assert cursor_path.read_text(encoding="utf-8").strip() == "c2"
    assert reader.stopped is True


@pytest.mark.asyncio
async def test_enricher_sets_hostname_and_static_fields() -> None:
    processor = EnricherProcessor(
        "host",
        config={"hostname": True, "static": {"env": "prod"}, "target": "metadata"},
    )
    record = Record(source="s", record_type="t", body={"message": "m"})
    result = await processor.process(record)
    assert result is not None
    assert result.metadata["hostname"]
    assert result.metadata["env"] == "prod"


@pytest.mark.asyncio
async def test_filter_drop_and_keep() -> None:
    dropper = FilterProcessor("drop", config={"drop": {"body.level": "debug"}})
    debug = Record(source="s", record_type="t", body={"level": "debug"})
    info = Record(source="s", record_type="t", body={"level": "info"})
    assert await dropper.process(debug) is None
    assert await dropper.process(info) is not None

    keeper = FilterProcessor("keep", config={"keep": {"source": "journal"}})
    assert await keeper.process(Record(source="other", record_type="t", body={})) is None
    kept = await keeper.process(Record(source="journal", record_type="t", body={}))
    assert kept is not None


@pytest.mark.asyncio
async def test_loki_sink_posts_stream() -> None:
    response = MagicMock()
    response.is_success = True
    response.status_code = 204
    client = AsyncMock()
    client.post = AsyncMock(return_value=response)
    client.aclose = AsyncMock()
    with patch("daena.pipeline.sinks.loki.httpx.AsyncClient", return_value=client):
        sink = LokiSink("loki", config={"url": "http://loki/push", "labels": {"job": "daena"}})
        await sink.start()
        record = Record(
            source="system",
            record_type="journald",
            body={"message": "hello"},
            timestamp="2020-09-13T12:26:40+00:00",
        )
        result = await sink.deliver(record)
        await sink.stop()
    assert result.success is True
    payload = client.post.await_args.kwargs["json"]
    stream = payload["streams"][0]
    assert stream["stream"]["job"] == "daena"
    assert stream["stream"]["source"] == "system"
    assert stream["values"][0][1] == "hello"
    assert stream["values"][0][0].isdigit()
    json.dumps(payload)