diff --git a/scripts/vedastro_gateway.py b/scripts/vedastro_gateway.py index 7daeed36..382a71b0 100644 --- a/scripts/vedastro_gateway.py +++ b/scripts/vedastro_gateway.py @@ -133,7 +133,7 @@ def _official_raw_response(result: dict[str, Any]) -> Any: if raw: return raw raw = result.get("raw_response") - if isinstance(raw, dict) and raw.get("source") == "vedastro_official": + if isinstance(raw, dict) and str(raw.get("source") or "").startswith("vedastro_official"): return raw return None diff --git a/tests/test_vedastro_gateway.py b/tests/test_vedastro_gateway.py index a136debb..5c27498c 100644 --- a/tests/test_vedastro_gateway.py +++ b/tests/test_vedastro_gateway.py @@ -261,6 +261,19 @@ def test_gateway_completion_archives_official_raw_response(monkeypatch, tmp_path assert '"vedastro_official"' in path.read_text(encoding="utf-8") +def test_gateway_archives_official_full_snapshot_raw_response(monkeypatch, tmp_path): + from scripts import vedastro_gateway + + monkeypatch.setenv("VEDASTRO_GATEWAY_QUEUE_DIR", str(tmp_path)) + job = vedastro_gateway.enqueue_gateway_job({"year": 1990}, question="public validation") + completed = vedastro_gateway.complete_gateway_job( + job["job_id"], + {"raw_response": {"source": "vedastro_official_full_snapshot", "response_hash": "public-demo"}}, + ) + + assert completed["raw_response_archive"]["official_raw_response_available"] is True + + def test_gateway_lists_official_raw_response_archives(monkeypatch, tmp_path): from scripts import vedastro_gateway