Files
aios/tests/test_worker_availability_guard.py
T

200 lines
7.0 KiB
Python

from __future__ import annotations
import json
from pathlib import Path
import sys
ROOT=Path(__file__).resolve().parents[1]
sys.path.insert(0,str(ROOT/"active-runtime/app"))
from aios_worker_availability_guard import AvailabilityGuard
class FakeSystemd:
def __init__(self, roots, now=1000.0):
self.now=now
self.actions=[]
self.units={}
for unit,root in roots.items():
self.units[unit]={
"LoadState":"loaded",
"ActiveState":"active",
"UnitFileState":"enabled",
"Environment":f"AIOS_WORKER_STATE_ROOT={root}",
}
self.units["aios-external-agent-executor.service"]={
"LoadState":"loaded",
"ActiveState":"active",
"UnitFileState":"enabled",
"Environment":"",
}
def property(self, unit, name):
return self.units.get(unit,{}).get(name,"")
def environment(self, unit):
raw=self.property(unit,"Environment")
result={}
for item in raw.split():
if "=" in item:
k,v=item.split("=",1)
result[k]=v
return result
def enable(self, unit):
self.actions.append((unit,"enable"))
self.units[unit]["UnitFileState"]="enabled"
def start(self, unit):
self.actions.append((unit,"start"))
self.units[unit]["ActiveState"]="active"
env=self.environment(unit)
if "AIOS_WORKER_STATE_ROOT" not in env:
return
root=Path(env["AIOS_WORKER_STATE_ROOT"])
root.mkdir(parents=True,exist_ok=True)
(root/"worker-liveness.json").write_text(json.dumps({
"schema":"aios.online-agent.worker-liveness.v1",
"agent_id":"agent-"+unit,
"state":"POLL_OK",
"last_seen_epoch":self.now,
}),encoding="utf-8")
def restart(self, unit):
self.actions.append((unit,"restart"))
self.units[unit]["ActiveState"]="active"
env=self.environment(unit)
root=Path(env["AIOS_WORKER_STATE_ROOT"])
root.mkdir(parents=True,exist_ok=True)
(root/"worker-liveness.json").write_text(json.dumps({
"schema":"aios.online-agent.worker-liveness.v1",
"agent_id":"agent-"+unit,
"state":"POLL_OK",
"last_seen_epoch":self.now,
}),encoding="utf-8")
def make_guard(tmp_path, fake):
worker_units=tuple(
unit for unit in fake.units
if unit!="aios-external-agent-executor.service"
)
return AvailabilityGuard(
systemd=fake,
worker_units=worker_units,
dependency_units=("aios-external-agent-executor.service",),
status_path=tmp_path/"run/status.json",
recovery_path=tmp_path/"run/recovery.json",
inhibit_path=tmp_path/"run/maintenance.lock",
max_stale_seconds=90,
restart_cooldown_seconds=60,
startup_wait_seconds=0,
now_fn=lambda:fake.now,
sleep_fn=lambda _seconds:None,
)
def write_liveness(root, now, *, state="POLL_OK"):
root.mkdir(parents=True,exist_ok=True)
(root/"worker-liveness.json").write_text(json.dumps({
"schema":"aios.online-agent.worker-liveness.v1",
"agent_id":"agent",
"state":state,
"last_seen_epoch":now,
}),encoding="utf-8")
def test_healthy_workers_need_no_action(tmp_path):
roots={
"generic.service":tmp_path/"generic",
"hermes.service":tmp_path/"hermes",
"openclaw.service":tmp_path/"openclaw",
}
fake=FakeSystemd(roots)
for root in roots.values():
write_liveness(root,fake.now)
result=make_guard(tmp_path,fake).run(heal=True)
assert result["status"]=="HEALTHY"
assert result["all_fresh"] is True
assert fake.actions==[]
def test_inactive_worker_is_started_and_becomes_fresh(tmp_path):
roots={"generic.service":tmp_path/"generic"}
fake=FakeSystemd(roots)
fake.units["generic.service"]["ActiveState"]="inactive"
result=make_guard(tmp_path,fake).run(heal=True)
assert ("generic.service","start") in fake.actions
assert result["all_fresh"] is True
def test_disabled_worker_is_enabled_and_started(tmp_path):
roots={"generic.service":tmp_path/"generic"}
fake=FakeSystemd(roots)
fake.units["generic.service"]["UnitFileState"]="disabled"
fake.units["generic.service"]["ActiveState"]="inactive"
result=make_guard(tmp_path,fake).run(heal=True)
assert ("generic.service","enable") in fake.actions
assert ("generic.service","start") in fake.actions
assert result["all_fresh"] is True
def test_active_but_stale_worker_is_restarted(tmp_path):
roots={"generic.service":tmp_path/"generic"}
fake=FakeSystemd(roots)
write_liveness(roots["generic.service"],fake.now-500)
result=make_guard(tmp_path,fake).run(heal=True)
assert ("generic.service","restart") in fake.actions
assert result["all_fresh"] is True
assert result["workers"][0]["action"]=="RESTART_STALE"
def test_restart_cooldown_prevents_restart_storm(tmp_path):
roots={"generic.service":tmp_path/"generic"}
fake=FakeSystemd(roots)
write_liveness(roots["generic.service"],fake.now-500)
guard=make_guard(tmp_path,fake)
guard.recovery_path.parent.mkdir(parents=True,exist_ok=True)
guard.recovery_path.write_text(json.dumps({
"schema":"aios.worker-availability.recovery.v1",
"last_restart_epoch":{"generic.service":fake.now-10},
}),encoding="utf-8")
result=guard.run(heal=True)
assert ("generic.service","restart") not in fake.actions
assert result["status"]=="DEGRADED"
assert result["workers"][0]["detail"]=="RESTART_COOLDOWN_ACTIVE"
def test_maintenance_inhibit_disables_healing(tmp_path):
roots={"generic.service":tmp_path/"generic"}
fake=FakeSystemd(roots)
fake.units["generic.service"]["ActiveState"]="inactive"
guard=make_guard(tmp_path,fake)
guard.inhibit_path.parent.mkdir(parents=True,exist_ok=True)
guard.inhibit_path.write_text("planned maintenance\n")
result=guard.run(heal=True)
assert result["status"]=="MAINTENANCE_INHIBITED"
assert fake.actions==[]
def test_dependency_is_enabled_and_started(tmp_path):
roots={"generic.service":tmp_path/"generic"}
fake=FakeSystemd(roots)
write_liveness(roots["generic.service"],fake.now)
fake.units["aios-external-agent-executor.service"]["UnitFileState"]="disabled"
fake.units["aios-external-agent-executor.service"]["ActiveState"]="inactive"
result=make_guard(tmp_path,fake).run(heal=True)
assert ("aios-external-agent-executor.service","enable") in fake.actions
assert ("aios-external-agent-executor.service","start") in fake.actions
assert result["all_fresh"] is True
def test_check_reports_stale_without_mutation(tmp_path):
roots={"generic.service":tmp_path/"generic"}
fake=FakeSystemd(roots)
write_liveness(roots["generic.service"],fake.now-500)
result=make_guard(tmp_path,fake).run(heal=False)
assert result["status"]=="DEGRADED"
assert result["all_fresh"] is False
assert fake.actions==[]