ebc8ad26b6
- Updated the job structure to include `anmod_id`, allowing for multiple mathematical models to be processed for each field and disease combination. - Enhanced the batch processing logic to discover and execute jobs for `(field, disease, anmod_id)` instead of just `(field, disease)`. - Modified the output file naming convention to include `anmod_id`, ensuring unique output for each model. - Adjusted database interactions to accommodate the new primary key structure based on `pathogen_model_id`. - Updated relevant documentation in `README.md` to reflect changes in job discovery, execution, and output generation. - Enhanced unit tests to validate the new functionality and ensure correct behavior across various scenarios.
271 lines
9.7 KiB
Python
271 lines
9.7 KiB
Python
"""Unit tests for the 09:00 readiness target: late jobs still run."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import tempfile
|
|
import unittest
|
|
from datetime import date, datetime, timedelta
|
|
from datetime import time as dtime
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
from pipeline.batch import _deadline_notice, run_batch
|
|
from pipeline.config import (
|
|
ConcurrencySettings,
|
|
LlmSettings,
|
|
ObservabilitySettings,
|
|
RetrySettings,
|
|
ScheduleSettings,
|
|
Settings,
|
|
SqlSettings,
|
|
WeaviateSettings,
|
|
WorklistSettings,
|
|
)
|
|
from pipeline.job import JobResult, output_stem
|
|
from pipeline.output import write_run_report
|
|
from pipeline.worklist import Job
|
|
|
|
|
|
def _at(hour: int, minute: int = 0, second: int = 0) -> datetime:
|
|
return datetime.combine(date.today(), dtime(hour, minute, second))
|
|
|
|
|
|
def _job(field_id: int = 4012, anmod_id: int = 10) -> Job:
|
|
return Job(
|
|
field_id=field_id,
|
|
crop_english="grapevine",
|
|
crop_italian="Vite",
|
|
model_name="PERONOSPORA",
|
|
disease_english="downy mildew",
|
|
organic=False,
|
|
station=7,
|
|
anmod_id=anmod_id,
|
|
)
|
|
|
|
|
|
def _settings() -> Settings:
|
|
llm = LlmSettings(
|
|
provider="gemini",
|
|
openai_model="m",
|
|
anthropic_model="m",
|
|
gemini_model="m",
|
|
max_tokens=8,
|
|
temperature=0.0,
|
|
thinking_level=None,
|
|
)
|
|
return Settings(
|
|
root=Path("."),
|
|
field_id=None,
|
|
disease_name=None,
|
|
worklist=WorklistSettings(field_allowlist=()),
|
|
concurrency=ConcurrencySettings(
|
|
workers=1,
|
|
sql_limit=1,
|
|
gemini_limit=1,
|
|
gemini_requests_per_minute=60,
|
|
),
|
|
retry=RetrySettings(attempts=1, initial_backoff_seconds=0.0, max_backoff_seconds=0.0),
|
|
schedule=ScheduleSettings(deadline="09:00"),
|
|
observability=ObservabilitySettings(
|
|
enabled=False,
|
|
project_name="p",
|
|
endpoint="http://localhost",
|
|
hide_prompts=True,
|
|
max_attribute_chars=100,
|
|
),
|
|
crops_vocab=Path("crops.yaml"),
|
|
diseases_vocab=Path("diseases.yaml"),
|
|
llm_query_synthesis=llm,
|
|
llm_advice_generation=llm,
|
|
sql=SqlSettings(
|
|
driver="d",
|
|
server="localhost",
|
|
database="db",
|
|
app_database="db",
|
|
username="u",
|
|
password="p",
|
|
),
|
|
weaviate=WeaviateSettings(host="localhost", http_port=8080, grpc_port=50051),
|
|
gemini_api_key="",
|
|
openai_api_key="",
|
|
anthropic_api_key="",
|
|
prompts_dir=Path("prompts"),
|
|
output_dir=Path("output"),
|
|
)
|
|
|
|
|
|
class DeadlineNoticeTests(unittest.TestCase):
|
|
def test_before_80_percent_is_quiet(self) -> None:
|
|
start = datetime(2026, 10, 14, 7, 0)
|
|
deadline = datetime(2026, 10, 14, 9, 0)
|
|
now = start + timedelta(seconds=0.8 * 7200 - 1)
|
|
self.assertIsNone(_deadline_notice(now, deadline, start))
|
|
|
|
def test_at_80_percent_and_still_before_deadline(self) -> None:
|
|
start = datetime(2026, 10, 14, 7, 0)
|
|
deadline = datetime(2026, 10, 14, 9, 0)
|
|
now = start + timedelta(seconds=0.8 * 7200)
|
|
self.assertLess(now, deadline)
|
|
self.assertEqual(_deadline_notice(now, deadline, start), "budget_80")
|
|
|
|
def test_exactly_at_deadline_is_late(self) -> None:
|
|
start = datetime(2026, 10, 14, 7, 0)
|
|
deadline = datetime(2026, 10, 14, 9, 0)
|
|
self.assertEqual(_deadline_notice(deadline, deadline, start), "late")
|
|
|
|
def test_after_deadline_is_late(self) -> None:
|
|
start = datetime(2026, 10, 14, 7, 0)
|
|
deadline = datetime(2026, 10, 14, 9, 0)
|
|
now = datetime(2026, 10, 14, 9, 5)
|
|
self.assertEqual(_deadline_notice(now, deadline, start), "late")
|
|
|
|
def test_batch_already_started_past_deadline_is_late(self) -> None:
|
|
deadline = datetime(2026, 10, 14, 9, 0)
|
|
start = datetime(2026, 10, 14, 10, 0)
|
|
now = datetime(2026, 10, 14, 10, 1)
|
|
self.assertEqual(_deadline_notice(now, deadline, start), "late")
|
|
|
|
|
|
class RunBatchDeadlineTests(unittest.TestCase):
|
|
def _run(
|
|
self,
|
|
jobs: list[Job],
|
|
*,
|
|
now,
|
|
skip_existing: bool = False,
|
|
existing: set[int] | None = None,
|
|
log_level: str = "WARNING",
|
|
) -> dict:
|
|
captured: dict = {}
|
|
prompts = MagicMock()
|
|
prompts.validate.return_value = []
|
|
|
|
def _write_report(output_dir, as_of, results, *, jobs_started_after_deadline=0):
|
|
captured["results"] = list(results)
|
|
captured["late"] = jobs_started_after_deadline
|
|
return {"totals": {}}
|
|
|
|
def _run_job(resources, job, as_of, dry_run=False):
|
|
return JobResult(job=job, status="ok", payload={"output_stem": str(job.field_id)})
|
|
|
|
now_fn = now if callable(now) else (lambda tz=None: now)
|
|
with (
|
|
patch("pipeline.batch.PromptRegistry", return_value=prompts),
|
|
patch("pipeline.batch.Resources", return_value=MagicMock()),
|
|
patch("pipeline.batch.build_jobs", return_value=jobs),
|
|
patch("pipeline.batch.run_job", side_effect=_run_job) as run_job,
|
|
patch("pipeline.batch.write_job_output"),
|
|
patch("pipeline.batch.write_run_report", side_effect=_write_report),
|
|
patch("pipeline.batch._existing_advice_keys", return_value=existing or set()),
|
|
patch("pipeline.batch.datetime") as frozen,
|
|
self.assertLogs("pipeline.batch", level=log_level) as logs,
|
|
):
|
|
frozen.now.side_effect = now_fn
|
|
frozen.combine = datetime.combine
|
|
captured["code"] = run_batch(
|
|
_settings(),
|
|
MagicMock(),
|
|
MagicMock(),
|
|
date(2026, 10, 14),
|
|
skip_existing_advice=skip_existing,
|
|
)
|
|
captured["run_job"] = run_job
|
|
captured["logs"] = [record.getMessage() for record in logs.records]
|
|
return captured
|
|
|
|
def test_job_after_deadline_still_runs(self) -> None:
|
|
captured = self._run([_job()], now=_at(10, 0))
|
|
|
|
self.assertEqual(captured["run_job"].call_count, 1)
|
|
self.assertEqual(captured["late"], 1)
|
|
self.assertEqual(captured["code"], 0)
|
|
self.assertEqual([result.status for result in captured["results"]], ["ok"])
|
|
self.assertTrue(all(result.status != "skipped_deadline" for result in captured["results"]))
|
|
self.assertEqual(sum("has passed" in message for message in captured["logs"]), 1)
|
|
|
|
def test_two_late_jobs_warn_once(self) -> None:
|
|
jobs = [_job(4012), _job(4013)]
|
|
captured = self._run(jobs, now=_at(10, 0))
|
|
|
|
self.assertEqual(captured["run_job"].call_count, 2)
|
|
self.assertEqual(captured["late"], 2)
|
|
self.assertEqual(sum("has passed" in message for message in captured["logs"]), 1)
|
|
|
|
def test_existing_advice_after_deadline_is_not_a_late_start(self) -> None:
|
|
job = _job()
|
|
captured = self._run(
|
|
[job],
|
|
now=_at(10, 0),
|
|
skip_existing=True,
|
|
existing={job.anmod_id},
|
|
)
|
|
|
|
self.assertEqual(captured["run_job"].call_count, 0)
|
|
self.assertEqual(captured["late"], 0)
|
|
self.assertEqual([result.status for result in captured["results"]], ["skipped_existing"])
|
|
|
|
def test_budget_warning_is_logged_once(self) -> None:
|
|
start = _at(7, 0)
|
|
at_80 = start + timedelta(seconds=0.8 * 2 * 3600)
|
|
later = at_80 + timedelta(minutes=5)
|
|
times = iter([start, at_80, later])
|
|
captured = self._run([_job(4012), _job(4013)], now=lambda tz=None: next(times))
|
|
|
|
self.assertEqual(captured["run_job"].call_count, 2)
|
|
self.assertEqual(captured["late"], 0)
|
|
self.assertEqual(sum("80%" in message for message in captured["logs"]), 1)
|
|
self.assertFalse(any("has passed" in message for message in captured["logs"]))
|
|
|
|
def test_skip_existing_is_per_pathogen_model(self) -> None:
|
|
stored = _job(anmod_id=20)
|
|
other = _job(anmod_id=11)
|
|
captured = self._run(
|
|
[stored, other],
|
|
now=_at(8, 0),
|
|
skip_existing=True,
|
|
existing={stored.anmod_id},
|
|
log_level="INFO",
|
|
)
|
|
|
|
self.assertEqual(captured["run_job"].call_count, 1)
|
|
self.assertEqual(captured["run_job"].call_args.args[1].anmod_id, 11)
|
|
self.assertEqual(
|
|
sorted(result.status for result in captured["results"]),
|
|
["ok", "skipped_existing"],
|
|
)
|
|
|
|
|
|
class OutputStemTests(unittest.TestCase):
|
|
def test_each_model_has_its_own_file_stem(self) -> None:
|
|
stems = {
|
|
output_stem(_job(anmod_id=11), "downy mildew"),
|
|
output_stem(_job(anmod_id=20), "downy mildew"),
|
|
}
|
|
self.assertEqual(
|
|
stems,
|
|
{
|
|
"field4012__grapevine__downy_mildew__model11",
|
|
"field4012__grapevine__downy_mildew__model20",
|
|
},
|
|
)
|
|
|
|
|
|
class RunReportTests(unittest.TestCase):
|
|
def test_report_records_late_starts_without_skipped_deadline(self) -> None:
|
|
result = JobResult(job=_job(), status="ok")
|
|
with tempfile.TemporaryDirectory() as tmp:
|
|
report = write_run_report(
|
|
Path(tmp),
|
|
date(2026, 10, 14),
|
|
[result],
|
|
jobs_started_after_deadline=2,
|
|
)
|
|
empty = write_run_report(Path(tmp), date(2026, 10, 14), [])
|
|
|
|
self.assertEqual(report["jobs_started_after_deadline"], 2)
|
|
self.assertEqual(report["jobs"][0]["anmod_id"], 10)
|
|
self.assertNotIn("skipped_deadline", report["totals"])
|
|
self.assertIn("skipped_existing", report["totals"])
|
|
self.assertEqual(empty["jobs_started_after_deadline"], 0)
|
|
self.assertNotIn("skipped_deadline", empty["totals"])
|