test_pipeline_formal.py 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204
  1. from __future__ import annotations
  2. from uuid import UUID, uuid4
  3. import pytest
  4. from acquisition.domain import CandidateItem, MediaAsset
  5. from acquisition.runner import RunBatchResult
  6. from core.models import Post
  7. from decode_content.models import GateResult, ReadResult
  8. from decode_content.service import DecodeWorkflowOutput
  9. from pipeline import acquisition_runner, creation_pipeline
  10. from pipeline.acquisition_runner import AcquisitionStageResult
  11. from pipeline.decode_runner import run_decode_stage
  12. from pipeline.decode_runner import DecodeBatchResult
  13. from pipeline.dedupe import dedupe_candidate_items, item_dedupe_key, should_decode_item
  14. from pipeline.models import PipelineJob, PipelineRun
  15. class FakePipelineRepo:
  16. def __init__(self):
  17. self.jobs = []
  18. def create_pipeline_run(self, **kwargs):
  19. return PipelineRun(id=uuid4(), **kwargs)
  20. def save_pipeline_job(self, **kwargs):
  21. job = PipelineJob(id=uuid4(), **kwargs)
  22. self.jobs.append(job)
  23. return job
  24. def mark_job_status(self, job_id, **kwargs):
  25. job = next(row for row in self.jobs if row.id == job_id)
  26. updated = job.model_copy(update=kwargs)
  27. self.jobs.append(updated)
  28. return updated
  29. def get_resume_cursor(self, run_id):
  30. return None
  31. def test_dedupe_candidate_items_prefers_platform_item_id_then_url():
  32. first = CandidateItem(id=uuid4(), platform="xiaohongshu", platform_item_id="x1", canonical_url="https://a")
  33. dup = CandidateItem(id=uuid4(), platform="xiaohongshu", platform_item_id="x1", canonical_url="https://b")
  34. url_only = CandidateItem(id=uuid4(), platform="weixin", canonical_url="HTTPS://MP.TEST/A")
  35. assert item_dedupe_key(first) == "xiaohongshu:id:x1"
  36. assert item_dedupe_key(url_only) == "url:https://mp.test/a"
  37. assert dedupe_candidate_items([first, dup, url_only]) == [first, url_only]
  38. assert should_decode_item(first, decoded_item_ids={str(first.id)}).keep is False
  39. def test_acquisition_stage_wraps_existing_runner(monkeypatch):
  40. batch_id = uuid4()
  41. pipeline_run = PipelineRun(id=uuid4(), status="running", current_stage="search")
  42. pipeline_repo = FakePipelineRepo()
  43. def fake_run_batch(acquisition_repo, *, batch_id, settings, **kwargs):
  44. return RunBatchResult(run_id=uuid4(), total_jobs=1, done=1, partial=0, failed=0, skipped=0)
  45. monkeypatch.setattr(acquisition_runner, "run_batch", fake_run_batch)
  46. result = acquisition_runner.run_acquisition_stage(
  47. acquisition_repo=object(),
  48. batch_id=batch_id,
  49. settings=object(),
  50. pipeline_repo=pipeline_repo,
  51. pipeline_run=pipeline_run,
  52. )
  53. assert result.acquisition.done == 1
  54. assert pipeline_repo.jobs[0].stage == "search"
  55. assert pipeline_repo.jobs[-1].status == "done"
  56. def test_acquisition_stage_marks_pipeline_job_failed_and_reraises(monkeypatch):
  57. pipeline_run = PipelineRun(id=uuid4(), status="running", current_stage="search")
  58. pipeline_repo = FakePipelineRepo()
  59. def fake_run_batch(*args, **kwargs):
  60. raise RuntimeError("acquisition exploded")
  61. monkeypatch.setattr(acquisition_runner, "run_batch", fake_run_batch)
  62. with pytest.raises(RuntimeError, match="acquisition exploded"):
  63. acquisition_runner.run_acquisition_stage(
  64. acquisition_repo=object(),
  65. batch_id=uuid4(),
  66. settings=object(),
  67. pipeline_repo=pipeline_repo,
  68. pipeline_run=pipeline_run,
  69. )
  70. assert pipeline_repo.jobs[0].status == "running"
  71. assert pipeline_repo.jobs[-1].status == "failed"
  72. assert pipeline_repo.jobs[-1].error_message == "acquisition exploded"
  73. class FakeCandidateRepo:
  74. def __init__(self):
  75. self.item_id = uuid4()
  76. self.items = [
  77. CandidateItem(
  78. id=self.item_id,
  79. platform="xiaohongshu",
  80. platform_item_id="x1",
  81. canonical_url="https://xhs/1",
  82. title="脚本",
  83. status="candidate",
  84. ),
  85. CandidateItem(
  86. id=uuid4(),
  87. platform="xiaohongshu",
  88. platform_item_id="x1",
  89. title="重复",
  90. status="candidate",
  91. ),
  92. ]
  93. def list_creation_candidate_items(self, *, run_id: UUID | None = None, limit: int = 100):
  94. return self.items
  95. def list_media_assets_for_item(self, item_id: UUID):
  96. return [
  97. MediaAsset(
  98. id=uuid4(),
  99. item_id=item_id,
  100. media_type="image",
  101. cdn_url="https://cdn.test/1.jpg",
  102. position=1,
  103. status="done",
  104. )
  105. ]
  106. class FakeDecodeService:
  107. def __init__(self):
  108. self.posts: list[Post] = []
  109. def decode_post(self, *, item_id: UUID, post: Post):
  110. self.posts.append(post)
  111. return DecodeWorkflowOutput(
  112. item_id=item_id,
  113. read_result=ReadResult(text="ok"),
  114. gate_result=GateResult(passed=True),
  115. status="decoded",
  116. )
  117. def test_decode_stage_loads_candidate_media_and_dedupes():
  118. repo = FakeCandidateRepo()
  119. service = FakeDecodeService()
  120. result = run_decode_stage(candidate_repo=repo, decode_service=service)
  121. assert result.total == 1
  122. assert result.decoded == 1
  123. assert result.failed == 0
  124. assert service.posts[0].image_urls == ["https://cdn.test/1.jpg"]
  125. def test_creation_pipeline_orchestrates_acquisition_then_decode(monkeypatch):
  126. batch_id = uuid4()
  127. acquisition_run_id = uuid4()
  128. pipeline_repo = FakePipelineRepo()
  129. calls = {}
  130. def fake_acquisition_stage(**kwargs):
  131. calls["acquisition"] = kwargs
  132. return AcquisitionStageResult(
  133. pipeline_job=None,
  134. acquisition=RunBatchResult(
  135. run_id=acquisition_run_id,
  136. total_jobs=3,
  137. done=3,
  138. partial=0,
  139. failed=0,
  140. skipped=0,
  141. ),
  142. )
  143. def fake_decode_stage(**kwargs):
  144. calls["decode"] = kwargs
  145. return DecodeBatchResult(total=2, decoded=2, skipped=0, failed=0, outputs=[])
  146. monkeypatch.setattr(creation_pipeline, "run_acquisition_stage", fake_acquisition_stage)
  147. monkeypatch.setattr(creation_pipeline, "run_decode_stage", fake_decode_stage)
  148. result = creation_pipeline.run_creation_pipeline(
  149. acquisition_repo=object(),
  150. batch_id=batch_id,
  151. settings=object(),
  152. pipeline_repo=pipeline_repo,
  153. decode_service=object(),
  154. decode_limit=20,
  155. )
  156. assert result.pipeline_run is not None
  157. assert result.pipeline_run.batch_id == batch_id
  158. assert result.acquisition.acquisition.done == 3
  159. assert result.decode is not None
  160. assert result.decode.decoded == 2
  161. assert calls["acquisition"]["pipeline_run"] == result.pipeline_run
  162. assert calls["decode"]["run_id"] == acquisition_run_id
  163. assert calls["decode"]["limit"] == 20