test_ingest_client.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135
  1. from __future__ import annotations
  2. from uuid import uuid4
  3. from core.config import IngestApiConfig
  4. from decode_content.ingest import KnowledgeIngestClient, ingest_payload_draft
  5. from decode_content.models import IngestRecord, PayloadDraft
  6. class FakeRepo:
  7. def __init__(self):
  8. self.marked = []
  9. self.records = []
  10. def mark_payload_draft_ingested(self, payload_draft_id):
  11. self.marked.append(payload_draft_id)
  12. return PayloadDraft(
  13. id=payload_draft_id,
  14. payload=_payload(),
  15. review_status="approved",
  16. ingest_ready=True,
  17. status="ingested",
  18. )
  19. def save_ingest_record(self, **kwargs):
  20. record = IngestRecord(id=uuid4(), **kwargs)
  21. self.records.append(record)
  22. return record
  23. class FakeResponse:
  24. def __init__(self, status_code, payload):
  25. self.status_code = status_code
  26. self._payload = payload
  27. self.text = str(payload)
  28. def json(self):
  29. return self._payload
  30. def _payload():
  31. return {
  32. "source": {"id": "xhs_1", "source_type": "post", "title": "帖子", "author": "作者"},
  33. "title": "拍照姿势",
  34. "content": "拍照时用手部动作制造自然感。",
  35. "dim_attributes": ["how"],
  36. "dim_creations": ["创作"],
  37. "scopes": [{"scope_type": "intent", "value": "自然感"}],
  38. "custom_ext": [{"key": "业务阶段", "type": "str", "value": "脚本"}],
  39. }
  40. def test_dry_run_ingest_validates_and_does_not_call_http():
  41. repo = FakeRepo()
  42. draft_id = uuid4()
  43. draft = PayloadDraft(id=draft_id, payload=_payload(), review_status="pending", ingest_ready=False, status="draft")
  44. record = ingest_payload_draft(repo, draft, dry_run=True)
  45. assert repo.marked == [draft_id]
  46. assert record.target_system == "dry-run"
  47. assert record.status == "ingested"
  48. assert record.response_payload["dry_run"] is True
  49. def test_real_ingest_posts_and_records_knowledge_id():
  50. calls = []
  51. def fake_post(url, **kwargs):
  52. calls.append((url, kwargs))
  53. return FakeResponse(201, {"knowledge_id": "k-1", "source_id": "xhs_1"})
  54. repo = FakeRepo()
  55. draft_id = uuid4()
  56. draft = PayloadDraft(id=draft_id, payload=_payload(), review_status="pending", ingest_ready=False, status="draft")
  57. client = KnowledgeIngestClient(
  58. IngestApiConfig(url="https://ingest.test/api/v1/knowledge/ingest", retry_delays_seconds=()),
  59. post_fn=fake_post,
  60. sleep_fn=lambda _delay: None,
  61. )
  62. record = ingest_payload_draft(repo, draft, dry_run=False, client=client)
  63. assert calls and calls[0][0] == "https://ingest.test/api/v1/knowledge/ingest"
  64. assert record.target_system == "knowledge-ingest-api"
  65. assert record.target_id == "k-1"
  66. assert record.response_payload["source_id"] == "xhs_1"
  67. assert repo.marked == [draft_id]
  68. def test_real_ingest_failure_is_recorded_without_marking_draft():
  69. def fake_post(url, **kwargs):
  70. return FakeResponse(502, {"detail": "embedding down"})
  71. repo = FakeRepo()
  72. draft = PayloadDraft(id=uuid4(), payload=_payload(), review_status="pending", ingest_ready=False, status="draft")
  73. client = KnowledgeIngestClient(
  74. IngestApiConfig(url="https://ingest.test/api/v1/knowledge/ingest", retry_delays_seconds=()),
  75. post_fn=fake_post,
  76. sleep_fn=lambda _delay: None,
  77. )
  78. record = ingest_payload_draft(repo, draft, dry_run=False, client=client)
  79. assert record.status == "failed"
  80. assert record.target_system == "knowledge-ingest-api"
  81. assert record.response_payload["http_status"] == 502
  82. assert repo.marked == []
  83. def test_real_ingest_retries_before_success():
  84. calls = []
  85. sleeps = []
  86. def fake_post(url, **kwargs):
  87. calls.append(url)
  88. if len(calls) == 1:
  89. return FakeResponse(502, {"detail": "embedding down"})
  90. return FakeResponse(201, {"knowledge_id": "k-2", "source_id": "xhs_1"})
  91. repo = FakeRepo()
  92. draft = PayloadDraft(id=uuid4(), payload=_payload(), review_status="pending", ingest_ready=False, status="draft")
  93. client = KnowledgeIngestClient(
  94. IngestApiConfig(url="https://ingest.test/api/v1/knowledge/ingest", retry_delays_seconds=(0.01,)),
  95. post_fn=fake_post,
  96. sleep_fn=lambda delay: sleeps.append(delay),
  97. )
  98. record = ingest_payload_draft(repo, draft, dry_run=False, client=client)
  99. assert len(calls) == 2
  100. assert sleeps == [0.01]
  101. assert record.status == "ingested"
  102. assert record.attempt_count == 2
  103. assert record.target_id == "k-2"