89 lines
3.4 KiB
Python
89 lines
3.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_outputs
|
|
|
|
|
|
class SoftwareOutputPublishTests(unittest.TestCase):
|
|
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()
|