112 lines
4.5 KiB
Python
112 lines
4.5 KiB
Python
import pytest
|
|
|
|
from app.agent import TOOL_SCHEMAS, ToolArgumentError, ToolExecutor, validate_arguments
|
|
from app.checkmk import Downtime, ServiceState
|
|
from app.poller.state import StateChange
|
|
|
|
|
|
def svc(host="web01", service="CPU load", state=2, output="CRIT", downtime=False):
|
|
return ServiceState(host, service, state, True, output, 1730000000, downtime, False)
|
|
|
|
|
|
class FakeReader:
|
|
def __init__(self):
|
|
self.calls = []
|
|
self.fail = False
|
|
self.items = [svc(), svc(service="Disk", state=1), svc(service="Mem", state=0, output="ok")]
|
|
|
|
async def service_status(self, host, service):
|
|
self.calls.append(("status", host, service))
|
|
if self.fail:
|
|
raise RuntimeError("secret-token boom")
|
|
return next((i for i in self.items if i.service == service), None)
|
|
|
|
async def host_services(self, host):
|
|
self.calls.append(("services", host))
|
|
return self.items
|
|
|
|
async def downtimes(self, host=None):
|
|
self.calls.append(("downtimes", host))
|
|
return [Downtime("web01", None, "2026-01-01", "2026-01-02", "Wartung", "admin")]
|
|
|
|
async def problems(self):
|
|
return self.items[:2]
|
|
|
|
|
|
def make_executor(reader=None, log=None):
|
|
return ToolExecutor(reader or FakeReader(), log or (lambda h, s: [StateChange(1730000000, 2)]))
|
|
|
|
|
|
def test_schemas_are_five_read_tools():
|
|
names = [s["function"]["name"] for s in TOOL_SCHEMAS]
|
|
assert names == ["get_service_status", "get_host_services", "get_state_history", "get_downtimes",
|
|
"get_problems"]
|
|
|
|
|
|
@pytest.mark.parametrize("tool,args", [
|
|
("get_service_status", {"host": "web01"}),
|
|
("get_service_status", {"host": "web 01", "service": "x"}),
|
|
("get_service_status", {"host": "web01", "service": "x\ny"}),
|
|
("get_service_status", {"host": "web01", "service": "x" * 201}),
|
|
("get_service_status", {"host": 5, "service": "x"}),
|
|
("get_service_status", {"host": "web01", "service": "x", "extra": "1"}),
|
|
("get_host_services", {"host": "../etc"}),
|
|
("get_host_services", {"host": "a" * 129}),
|
|
("get_problems", {"host": "x"}),
|
|
("get_downtimes", {"host": "a;b"}),
|
|
("get_host_services", "web01"),
|
|
("rm_rf", {}),
|
|
])
|
|
def test_invalid_arguments(tool, args):
|
|
with pytest.raises(ToolArgumentError):
|
|
validate_arguments(tool, args)
|
|
|
|
|
|
def test_valid_arguments():
|
|
assert validate_arguments("get_service_status", {"host": "web01.example", "service": " CPU load "}) == \
|
|
{"host": "web01.example", "service": "CPU load"}
|
|
assert validate_arguments("get_downtimes", {}) == {}
|
|
assert validate_arguments("get_downtimes", {"host": None}) == {}
|
|
assert validate_arguments("get_downtimes", {"host": "web01"}) == {"host": "web01"}
|
|
assert validate_arguments("get_service_status", {"host": "h", "service": "Filesystem /var (data)"})
|
|
|
|
|
|
async def test_executor_each_tool():
|
|
ex = make_executor()
|
|
text, err = await ex.execute("get_service_status", {"host": "web01", "service": "CPU load"})
|
|
assert err is None and "CRIT web01 / CPU load" in text
|
|
text, _ = await ex.execute("get_service_status", {"host": "web01", "service": "nope"})
|
|
assert "nicht gefunden" in text
|
|
text, _ = await ex.execute("get_host_services", {"host": "web01"})
|
|
assert text.splitlines()[0].startswith("CRIT")
|
|
text, _ = await ex.execute("get_state_history", {"host": "web01", "service": "CPU load"})
|
|
assert "CRIT" in text
|
|
text, _ = await ex.execute("get_downtimes", {})
|
|
assert "Wartung" in text
|
|
text, _ = await ex.execute("get_problems", {})
|
|
assert "Disk" in text and "Mem" not in text
|
|
|
|
|
|
async def test_async_state_log():
|
|
async def log(host, service):
|
|
return [StateChange(1730000000, 1)]
|
|
text, err = await make_executor(log=log).execute("get_state_history", {"host": "h", "service": "s"})
|
|
assert err is None and "WARN" in text
|
|
|
|
|
|
async def test_errors_never_raise_and_do_not_leak():
|
|
reader = FakeReader()
|
|
reader.fail = True
|
|
text, err = await make_executor(reader).execute("get_service_status", {"host": "h", "service": "s"})
|
|
assert err == "RuntimeError" and "secret-token" not in text
|
|
text, err = await make_executor(reader).execute("nope", {})
|
|
assert err and text.startswith("Fehler")
|
|
assert reader.calls == [("status", "h", "s")]
|
|
|
|
|
|
async def test_result_truncated_and_list_limited():
|
|
reader = FakeReader()
|
|
reader.items = [svc(service=f"s{i}", output="x" * 500) for i in range(100)]
|
|
text, _ = await make_executor(reader).execute("get_host_services", {"host": "h"})
|
|
assert len(text) <= 1500
|