205 lines
8.4 KiB
Python
205 lines
8.4 KiB
Python
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_01": ("preview-01.png", b"overview"),
|
|
"preview_02": ("preview-02.png", b"front"),
|
|
}
|
|
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@v3",
|
|
"user_id": uuid4(),
|
|
"working_dir": str(root / "research"),
|
|
"request": {
|
|
"outputs": [
|
|
{"key": "preview_02", "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_01", "preview_02"},
|
|
)
|
|
|
|
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()
|