from __future__ import annotations import hashlib import tempfile import unittest from pathlib import Path from unittest.mock import patch from uuid import uuid4 from web.routers.software_nodes import ( _publish_software_job_export, _publish_software_job_outputs, _publish_software_job_previews, ) class SoftwareOutputPublishTests(unittest.TestCase): def test_workspace_preview_returns_to_task_folder_without_artifact_registration(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) job_id = uuid4() staging = root / ".zcbot_software_job_staging" / str(job_id) staging.mkdir(parents=True) content = b"preview" (staging / "figure.png").write_bytes(content) manifest = [{ "artifact_id": "figure_png", "filename": "figure.png", "media_type": "image/png", "size_bytes": len(content), "sha256": hashlib.sha256(content).hexdigest(), }] context = { "capability": "origin.plot@v2", "user_id": uuid4(), "working_dir": str(root / "research"), } (root / "research").mkdir() with ( patch("web.routers.software_nodes.load_user_root", return_value=root), patch("web.routers.software_nodes.register_published_artifacts") as register, ): previews = _publish_software_job_previews(job_id, context, manifest) self.assertFalse(register.called) self.assertNotIn("artifact_id", previews[0]) self.assertEqual(previews[0]["preview_id"], "figure_png") self.assertEqual(previews[0]["path"], f"origin/{job_id}/figure.png") self.assertEqual( (root / "research" / "origin" / str(job_id) / "figure.png").read_bytes(), content, ) def test_workspace_publishes_optional_previews_requested_by_the_job(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) job_id = uuid4() staging = root / ".zcbot_software_job_staging" / str(job_id) staging.mkdir(parents=True) files = { "preview_isometric": ("preview-isometric.png", b"isometric"), "preview_cutaway": ("preview-cutaway.png", b"cutaway"), } manifest = [] for output_id, (filename, content) in files.items(): (staging / filename).write_bytes(content) manifest.append({ "artifact_id": output_id, "filename": filename, "media_type": "image/png", "size_bytes": len(content), "sha256": hashlib.sha256(content).hexdigest(), }) context = { "capability": "blender.scene.author@v2", "user_id": uuid4(), "working_dir": str(root / "research"), "request": { "outputs": [ {"key": "preview_cutaway", "type": "figure", "format": "png"} ] }, } (root / "research").mkdir() with patch("web.routers.software_nodes.load_user_root", return_value=root): previews = _publish_software_job_previews(job_id, context, manifest) self.assertEqual( {item["preview_id"] for item in previews}, {"preview_isometric", "preview_cutaway"}, ) def test_explicit_export_registers_an_existing_preview_file(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) job_id = uuid4() working_dir = root / "research" preview = working_dir / "origin" / str(job_id) / "figure.png" preview.parent.mkdir(parents=True) content = b"preview" preview.write_bytes(content) manifest = [{ "artifact_id": "figure_png", "filename": "figure.png", "media_type": "image/png", "size_bytes": len(content), "sha256": hashlib.sha256(content).hexdigest(), }] context = { "capability": "origin.plot@v2", "user_id": uuid4(), "task_id": uuid4(), "working_dir": str(working_dir), } published_id = str(uuid4()) with ( patch("web.routers.software_nodes.load_user_root", return_value=root), patch( "web.routers.software_nodes.register_published_artifacts", return_value=({ "path": f"origin/{job_id}/figure.png", "artifact_id": published_id, },), ) as register, ): published = _publish_software_job_export(job_id, context, manifest) self.assertEqual(preview.read_bytes(), content) self.assertEqual(published[0]["artifact_id"], published_id) self.assertEqual(register.call_args.kwargs["software_job_id"], job_id) def test_complete_set_moves_atomically_and_can_be_replayed(self) -> None: with tempfile.TemporaryDirectory() as directory: root = Path(directory) job_id = uuid4() working_dir = root / "research" staging = root / ".zcbot_software_job_staging" / str(job_id) staging.mkdir(parents=True) working_dir.mkdir() content = b"origin-result" metadata = b'{"kind":"plot"}' (staging / "figure.png").write_bytes(content) (staging / "plot-spec.json").write_bytes(metadata) manifest = [ { "artifact_id": "figure_png", "filename": "figure.png", "media_type": "image/png", "size_bytes": len(content), "sha256": hashlib.sha256(content).hexdigest(), }, { "artifact_id": "plot_spec", "filename": "plot-spec.json", "media_type": "application/json", "size_bytes": len(metadata), "sha256": hashlib.sha256(metadata).hexdigest(), }, ] context = { "user_id": uuid4(), "task_id": uuid4(), "working_dir": str(working_dir), } registered = [] def register(**kwargs): registered.append(kwargs) return tuple({ "version": 2, "scope": "working_dir", "path": ref["path"], "label": ref["label"], "artifact_id": str(uuid4()), } for ref in kwargs["refs"]) with ( patch("web.routers.software_nodes.load_user_root", return_value=root), patch( "web.routers.software_nodes.register_published_artifacts", side_effect=register, ), ): first = _publish_software_job_outputs(job_id, context, manifest) second = _publish_software_job_outputs(job_id, context, manifest) published = working_dir / "origin" / str(job_id) / "figure.png" self.assertEqual(published.read_bytes(), content) self.assertEqual( (working_dir / "origin" / str(job_id) / ".meta" / "plot-spec.json").read_bytes(), metadata, ) self.assertFalse(staging.exists()) self.assertEqual(first[0]["source_artifact_id"], "figure_png") self.assertEqual(first[0]["path"], f"origin/{job_id}/figure.png") self.assertEqual(first[1]["source_artifact_id"], "plot_spec") self.assertIsNone(first[1]["artifact_id"]) self.assertEqual(first[1]["path"], f"origin/{job_id}/.meta/plot-spec.json") self.assertEqual(second[0]["source_artifact_id"], "figure_png") self.assertEqual(registered[0]["software_job_id"], job_id) self.assertEqual(len(registered[0]["refs"]), 1) if __name__ == "__main__": unittest.main()