124 lines
3.4 KiB
Python
124 lines
3.4 KiB
Python
import threading
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import bot_descarga.worker.agent as agent_module
|
|
from bot_descarga.models.resultado import RunSummary
|
|
from bot_descarga.worker.agent import WorkerAgent
|
|
|
|
|
|
class FakeApi:
|
|
def __init__(self, failures: int = 0):
|
|
self.calls = 0
|
|
self.failures = failures
|
|
|
|
def heartbeat(self, worker_id: str) -> dict:
|
|
self.calls += 1
|
|
if self.calls <= self.failures:
|
|
raise RuntimeError("gateway temporalmente no disponible")
|
|
return {"worker_id": worker_id}
|
|
|
|
|
|
def make_agent(api: FakeApi) -> WorkerAgent:
|
|
agent = WorkerAgent.__new__(WorkerAgent)
|
|
agent.api = api
|
|
agent.worker_id = "zero0"
|
|
agent.heartbeat_seconds = 0.01
|
|
agent._log = lambda _message: None
|
|
return agent
|
|
|
|
|
|
def run_briefly(agent: WorkerAgent) -> threading.Thread:
|
|
stop_event = threading.Event()
|
|
thread = threading.Thread(
|
|
target=agent._heartbeat_while_job_runs,
|
|
args=(stop_event,),
|
|
daemon=True,
|
|
)
|
|
thread.start()
|
|
deadline = time.monotonic() + 1.0
|
|
while agent.api.calls < 3 and time.monotonic() < deadline:
|
|
time.sleep(0.01)
|
|
stop_event.set()
|
|
thread.join(timeout=1.0)
|
|
return thread
|
|
|
|
|
|
def test_active_job_sends_repeated_heartbeats() -> None:
|
|
api = FakeApi()
|
|
thread = run_briefly(make_agent(api))
|
|
|
|
assert api.calls >= 3
|
|
assert not thread.is_alive()
|
|
|
|
|
|
def test_transient_heartbeat_failure_does_not_stop_loop() -> None:
|
|
api = FakeApi(failures=1)
|
|
thread = run_briefly(make_agent(api))
|
|
|
|
assert api.calls >= 3
|
|
assert not thread.is_alive()
|
|
|
|
|
|
def test_successful_job_stays_completed_and_stops_heartbeat(monkeypatch, tmp_path: Path) -> None:
|
|
class JobApi(FakeApi):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.patches = []
|
|
|
|
def download_excel(self, _job_id: str, dest: Path) -> Path:
|
|
dest.parent.mkdir(parents=True, exist_ok=True)
|
|
dest.write_bytes(b"xlsx")
|
|
return dest
|
|
|
|
def cancel_flag(self, _job_id: str) -> bool:
|
|
return False
|
|
|
|
def post_event(self, _job_id: str, _event: dict) -> None:
|
|
return None
|
|
|
|
def patch_job(self, job_id: str, **body) -> dict:
|
|
self.patches.append((job_id, body))
|
|
return body
|
|
|
|
class FakeMinio:
|
|
def write_manifest(self, *_args, **_kwargs) -> None:
|
|
return None
|
|
|
|
class FakeDownloadService:
|
|
def __init__(self, **_kwargs):
|
|
pass
|
|
|
|
def run(self) -> RunSummary:
|
|
return RunSummary(total=1, exitosos=1)
|
|
|
|
api = JobApi()
|
|
agent = make_agent(api)
|
|
agent.staging_root = tmp_path
|
|
agent.minio = FakeMinio()
|
|
agent._uploaded_keys = []
|
|
monkeypatch.setattr(agent_module, "DownloadService", FakeDownloadService)
|
|
|
|
agent._run_job({"id": "job-1", "excel_key": "input.xlsx"})
|
|
calls_after_job = api.calls
|
|
time.sleep(0.03)
|
|
|
|
assert api.patches == [
|
|
(
|
|
"job-1",
|
|
{
|
|
"status": "completed",
|
|
"result": {
|
|
"total": 1,
|
|
"exitosos": 1,
|
|
"fallidos": 0,
|
|
"cancelled": False,
|
|
"downloads_prefix": "downloads/job-1/",
|
|
"uploaded_count": 0,
|
|
},
|
|
"error": None,
|
|
},
|
|
)
|
|
]
|
|
assert api.calls == calls_after_job
|