66 lines
2.4 KiB
Python
66 lines
2.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.compute_nodes import _publish_compute_outputs
|
|
|
|
|
|
class ComputeOutputPublishTests(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_compute_staging" / str(job_id)
|
|
staging.mkdir(parents=True)
|
|
working_dir.mkdir()
|
|
content = b"origin-result"
|
|
(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 = {
|
|
"user_id": uuid4(),
|
|
"task_id": uuid4(),
|
|
"working_dir": "research",
|
|
}
|
|
|
|
def register(**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.compute_nodes.load_user_root", return_value=root),
|
|
patch(
|
|
"web.routers.compute_nodes.register_published_artifacts",
|
|
side_effect=register,
|
|
),
|
|
):
|
|
first = _publish_compute_outputs(job_id, context, manifest)
|
|
second = _publish_compute_outputs(job_id, context, manifest)
|
|
|
|
published = working_dir / "origin" / str(job_id) / "figure.png"
|
|
self.assertEqual(published.read_bytes(), content)
|
|
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(second[0]["source_artifact_id"], "figure_png")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|