test_segment_entry.py 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236
  1. from __future__ import annotations
  2. import tempfile
  3. import unittest
  4. from pathlib import Path
  5. from unittest.mock import patch
  6. from production_build_agents.agents.segment.run import (
  7. run_segment_production,
  8. )
  9. from production_build_agents.run.records import VersionConflictError
  10. from production_build_agents.run.registry import register_thread
  11. from production_build_agents.production.runtime import (
  12. PRODUCTION_PROTOCOL_VERSION,
  13. validate_production_run_directory_boundary as validate_boundary,
  14. )
  15. from production_build_agents.run.layout import (
  16. run_directory,
  17. segment_package_path,
  18. )
  19. from production_build_agents.run.lock import (
  20. RunLockError,
  21. acquire_run_lock,
  22. )
  23. from run_segment import run_production_segment, segment_request_sha256
  24. from tests.support.production_fixtures import segment_bundle
  25. class SegmentEntryTest(unittest.TestCase):
  26. def test_thread_binds_exact_segment_package_bytes(self) -> None:
  27. with tempfile.TemporaryDirectory() as temp_dir:
  28. root = Path(temp_dir)
  29. package, *_ = segment_bundle(1, run_id="segment-thread")
  30. package_path = root / "Segment1.v1.json"
  31. package_path.write_text(
  32. package.model_dump_json(),
  33. encoding="utf-8",
  34. )
  35. registry = root / "registry"
  36. run_dir = root / "run"
  37. first_hash = segment_request_sha256(package_path.read_bytes())
  38. register_thread(
  39. registry_dir=registry,
  40. thread_id="segment-thread",
  41. protocol_version=PRODUCTION_PROTOCOL_VERSION,
  42. input_path=package_path,
  43. input_sha256=first_hash,
  44. run_dir=run_dir,
  45. )
  46. package_path.write_bytes(package_path.read_bytes() + b"\n")
  47. with self.assertRaisesRegex(
  48. VersionConflictError,
  49. "input_sha256",
  50. ):
  51. register_thread(
  52. registry_dir=registry,
  53. thread_id="segment-thread",
  54. protocol_version=PRODUCTION_PROTOCOL_VERSION,
  55. input_path=package_path,
  56. input_sha256=segment_request_sha256(
  57. package_path.read_bytes()
  58. ),
  59. run_dir=run_dir,
  60. )
  61. def test_old_run_is_rejected_before_registry_or_lock_write(self) -> None:
  62. with tempfile.TemporaryDirectory() as temp_dir:
  63. root = Path(temp_dir)
  64. package, *_ = segment_bundle(1, run_id="segment-thread")
  65. package_path = root / "Segment1.v1.json"
  66. package_path.write_text(
  67. package.model_dump_json(),
  68. encoding="utf-8",
  69. )
  70. output_root = root / "output"
  71. run_dir = run_directory(output_root, package.run_id)
  72. run_dir.mkdir(parents=True)
  73. old_record = run_dir / "production_package.json"
  74. old_record.write_text(
  75. '{"schema_version":"0.5"}',
  76. encoding="utf-8",
  77. )
  78. registry = root / "registry"
  79. with self.assertRaisesRegex(ValueError, "旧版正式路径"):
  80. run_production_segment(
  81. package_path,
  82. output_root=output_root,
  83. registry_dir=registry,
  84. )
  85. self.assertFalse(registry.exists())
  86. self.assertEqual(
  87. sorted(
  88. path.relative_to(run_dir)
  89. for path in run_dir.rglob("*")
  90. ),
  91. [Path("production_package.json")],
  92. )
  93. def test_locked_recheck_rejects_concurrent_mixed_record(self) -> None:
  94. with tempfile.TemporaryDirectory() as temp_dir:
  95. root = Path(temp_dir)
  96. package, *_ = segment_bundle(1, run_id="segment-thread")
  97. package_path = root / "Segment1.v1.json"
  98. package_path.write_text(
  99. package.model_dump_json(),
  100. encoding="utf-8",
  101. )
  102. output_root = root / "output"
  103. run_dir = run_directory(output_root, package.run_id)
  104. registry = root / "registry"
  105. injected = run_dir / "segment_packages" / "Injected.v1.json"
  106. call_count = 0
  107. def validate_then_inject(path: Path) -> None:
  108. nonlocal call_count
  109. validate_boundary(path)
  110. call_count += 1
  111. if call_count == 1:
  112. injected.parent.mkdir(parents=True)
  113. injected.write_text(
  114. '{"schema_version":"0.4"}',
  115. encoding="utf-8",
  116. )
  117. with (
  118. patch(
  119. "run_segment.validate_production_run_directory_boundary",
  120. side_effect=validate_then_inject,
  121. ),
  122. self.assertRaisesRegex(ValueError, "非当前协议"),
  123. ):
  124. run_production_segment(
  125. package_path,
  126. output_root=output_root,
  127. registry_dir=registry,
  128. )
  129. self.assertEqual(call_count, 1)
  130. self.assertFalse(registry.exists())
  131. self.assertEqual(
  132. injected.read_text(encoding="utf-8"),
  133. '{"schema_version":"0.4"}',
  134. )
  135. self.assertFalse(
  136. segment_package_path(
  137. run_dir,
  138. package.segment_id,
  139. package.plan_version,
  140. ).exists()
  141. )
  142. def test_source_mutation_after_parse_is_rejected_before_registration(
  143. self,
  144. ) -> None:
  145. with tempfile.TemporaryDirectory() as temp_dir:
  146. root = Path(temp_dir)
  147. package, *_ = segment_bundle(1, run_id="segment-thread")
  148. package_path = root / "Segment1.v1.json"
  149. package_path.write_text(
  150. package.model_dump_json(),
  151. encoding="utf-8",
  152. )
  153. original_bytes = package_path.read_bytes()
  154. output_root = root / "output"
  155. run_dir = run_directory(output_root, package.run_id)
  156. registry = root / "registry"
  157. call_count = 0
  158. def validate_then_mutate(path: Path) -> None:
  159. nonlocal call_count
  160. validate_boundary(path)
  161. call_count += 1
  162. if call_count == 1:
  163. package_path.write_bytes(original_bytes + b"\n")
  164. with (
  165. patch(
  166. "run_segment.validate_production_run_directory_boundary",
  167. side_effect=validate_then_mutate,
  168. ),
  169. self.assertRaisesRegex(
  170. VersionConflictError,
  171. "输入文件在封印前已变化",
  172. ),
  173. ):
  174. run_production_segment(
  175. package_path,
  176. output_root=output_root,
  177. registry_dir=registry,
  178. )
  179. self.assertEqual(call_count, 2)
  180. self.assertFalse(registry.exists())
  181. self.assertFalse(
  182. segment_package_path(
  183. run_dir,
  184. package.segment_id,
  185. package.plan_version,
  186. ).exists()
  187. )
  188. def test_locked_preflight_runs_while_run_lock_is_held(self) -> None:
  189. class PreflightObserved(RuntimeError):
  190. pass
  191. with tempfile.TemporaryDirectory() as temp_dir:
  192. root = Path(temp_dir)
  193. package, *_ = segment_bundle(1, run_id="segment-thread")
  194. def observe_lock(run_dir: Path) -> None:
  195. with self.assertRaises(RunLockError):
  196. with acquire_run_lock(run_dir):
  197. pass
  198. raise PreflightObserved
  199. with self.assertRaises(PreflightObserved):
  200. run_segment_production(
  201. package,
  202. run_dir=root / package.run_id,
  203. locked_preflight=observe_lock,
  204. )
  205. self.assertFalse(
  206. segment_package_path(
  207. root / package.run_id,
  208. package.segment_id,
  209. package.plan_version,
  210. ).exists()
  211. )
  212. if __name__ == "__main__":
  213. unittest.main()