daena/tests/test_grpc.py

from __future__ import annotations

import asyncio

import grpc  # noqa: I001
import pytest

from daena.api.grpc_server import DaenaGRPCServer
from daena.api.pb2 import daena_pb2 as pb2
from daena.api.pb2 import daena_pb2_grpc as pb2_grpc
from daena.config import DaenaConfig, PipelineConfig, SinkConfig, SourceConfig
from daena.core.runtime import DaenaRuntime
from daena.persistence.models import RecordState


@pytest.fixture
async def grpc_env(tmp_path):
    config = DaenaConfig()
    config.sources = [
        SourceConfig(name="dummy", type="dummy", options={"interval": 0.2, "message": "test"})
    ]
    config.sinks = [SinkConfig(name="log", type="dummy")]
    config.persistence.lmdb.path = str(tmp_path / "lmdb")
    config.pipeline = PipelineConfig(batch_size=100, flush_interval=5.0, max_queue_size=10_000)
    config.api.grpc.port = 18642

    runtime = DaenaRuntime(config)
    await runtime.start()

    server = DaenaGRPCServer(runtime, host="127.0.0.1", port=18642)
    await server.start()

    channel = grpc.aio.insecure_channel("127.0.0.1:18642")
    stub = pb2_grpc.DaenaStub(channel)

    yield runtime, stub, server

    await channel.close()
    await server.stop()
    await runtime.stop()


@pytest.mark.asyncio
async def test_get_status(grpc_env) -> None:
    _, stub, _ = grpc_env
    resp = await stub.GetStatus(pb2.Empty(), timeout=5)
    assert resp.service == "daena"
    assert resp.version == "0.2.0"
    assert resp.running is True


@pytest.mark.asyncio
async def test_get_health(grpc_env) -> None:
    _, stub, _ = grpc_env
    resp = await stub.GetHealth(pb2.Empty(), timeout=5)
    assert resp.healthy is True
    assert resp.running is True


@pytest.mark.asyncio
async def test_get_config(grpc_env) -> None:
    _, stub, _ = grpc_env
    resp = await stub.GetConfig(pb2.Empty(), timeout=5)
    assert resp.config_json != ""


@pytest.mark.asyncio
async def test_list_plugins(grpc_env) -> None:
    _, stub, _ = grpc_env
    resp = await stub.ListPlugins(pb2.Empty(), timeout=5)
    assert len(resp.sources.items) > 0
    assert len(resp.sinks.items) > 0


@pytest.mark.asyncio
async def test_get_queue_state(grpc_env) -> None:
    _, stub, _ = grpc_env
    await asyncio.sleep(1)
    resp = await stub.GetQueueState(pb2.Empty(), timeout=5)
    assert resp.total > 0
    assert resp.pending + resp.delivered > 0


@pytest.mark.asyncio
async def test_list_pending_records(grpc_env) -> None:
    runtime, stub, _ = grpc_env
    await asyncio.sleep(1)
    req = pb2.Pagination(limit=10, offset=0)
    resp = await stub.ListPendingRecords(req, timeout=5)
    assert resp.count > 0
    assert len(resp.records) > 0
    assert resp.records[0].record_id != ""


@pytest.mark.asyncio
async def test_list_sinks(grpc_env) -> None:
    _, stub, _ = grpc_env
    resp = await stub.ListSinks(pb2.Empty(), timeout=5)
    assert len(resp.sinks) >= 1
    assert resp.sinks[0].name == "log"


@pytest.mark.asyncio
async def test_replay_dead(grpc_env) -> None:
    runtime, stub, _ = grpc_env
    await asyncio.sleep(1)
    assert runtime.backend is not None
    stored = None
    for state in (RecordState.pending, RecordState.delivered, RecordState.failed):
        found = runtime.backend.list_by_state(state, limit=1)
        if found:
            stored = found[0]
            break
    assert stored is not None
    for sink in stored.deliveries:
        runtime.backend.dead(stored.record_id, "grpc replay", sink=sink)
    resp = await stub.ReplayDead(
        pb2.ReplayRequest(record_id=stored.record_id, limit=10),
        timeout=5,
    )
    assert resp.replayed == 1
    again = runtime.backend.get(stored.record_id)
    assert again is not None
    assert again.deliveries["log"].state == RecordState.pending


@pytest.mark.asyncio
async def test_reload_config(grpc_env) -> None:
    _, stub, _ = grpc_env
    resp = await stub.ReloadConfig(pb2.Empty(), timeout=10)
    assert resp.success is True