773 lines
34 KiB
Python
773 lines
34 KiB
Python
from __future__ import annotations
|
||
|
||
import importlib
|
||
import unittest
|
||
from datetime import datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
from unittest.mock import AsyncMock, MagicMock, patch
|
||
from uuid import uuid4
|
||
|
||
from alembic.migration import MigrationContext
|
||
from alembic.operations import Operations
|
||
from sqlalchemy import create_mock_engine
|
||
from sqlalchemy.dialects import postgresql
|
||
|
||
from core.software_jobs import (
|
||
SoftwareJobError,
|
||
_canonical_request,
|
||
_published_output_is_valid,
|
||
abandon_offer,
|
||
list_jobs,
|
||
mark_node_jobs_disconnected,
|
||
offer_next_job,
|
||
record_job_terminal,
|
||
replay_succeeded_outputs,
|
||
request_job_cancel,
|
||
respond_to_offer,
|
||
succeeded_output_upload_matches,
|
||
update_job_state,
|
||
validate_output_manifest,
|
||
)
|
||
from core.software_nodes import (
|
||
_enrollment_digest,
|
||
_hash_secret,
|
||
_verify_secret,
|
||
delete_node,
|
||
)
|
||
from web.routers.software_nodes import NodeConnectionManager, _bearer
|
||
|
||
|
||
class SoftwareNodeSecurityTests(unittest.TestCase):
|
||
def test_secret_hash_is_salted_and_verifiable(self) -> None:
|
||
first = _hash_secret("node-secret")
|
||
second = _hash_secret("node-secret")
|
||
self.assertNotEqual(first, second)
|
||
self.assertNotIn("node-secret", first)
|
||
self.assertTrue(_verify_secret("node-secret", first))
|
||
self.assertFalse(_verify_secret("wrong", first))
|
||
|
||
def test_bearer_parser_rejects_query_style_or_missing_token(self) -> None:
|
||
self.assertEqual(_bearer("Bearer abc"), "abc")
|
||
with self.assertRaisesRegex(Exception, "missing node bearer token"):
|
||
_bearer(None)
|
||
|
||
def test_enrollment_digest_does_not_store_plaintext(self) -> None:
|
||
digest = _enrollment_digest("ZCN-ABC")
|
||
self.assertEqual(len(digest), 64)
|
||
self.assertNotIn("ZCN-ABC", digest)
|
||
|
||
def test_websocket_auth_rejection_uses_explicit_application_close_code(self) -> None:
|
||
source = (
|
||
Path(__file__).resolve().parents[1] / "web" / "routers" / "software_nodes.py"
|
||
).read_text(encoding="utf-8")
|
||
rejection = source.split("except (ValueError, SoftwareNodeError):", 1)[1].split(
|
||
"await node_connections.activate", 1
|
||
)[0]
|
||
self.assertLess(
|
||
rejection.index("await websocket.accept()"), rejection.index("await websocket.close")
|
||
)
|
||
self.assertIn('code=4003, reason="invalid node credentials"', rejection)
|
||
|
||
|
||
class SoftwareNodeConnectionTests(unittest.IsolatedAsyncioTestCase):
|
||
async def test_new_connection_replaces_old_without_removing_new(self) -> None:
|
||
manager = NodeConnectionManager()
|
||
node_id = uuid4()
|
||
old = AsyncMock()
|
||
new = AsyncMock()
|
||
|
||
await manager.activate(node_id, old)
|
||
await manager.activate(node_id, new)
|
||
|
||
old.close.assert_awaited_once_with(
|
||
code=4001, reason="replaced by a newer connection"
|
||
)
|
||
self.assertFalse(await manager.remove(node_id, old))
|
||
self.assertTrue(await manager.remove(node_id, new))
|
||
|
||
async def test_admin_close_removes_and_closes_connection(self) -> None:
|
||
manager = NodeConnectionManager()
|
||
node_id = uuid4()
|
||
websocket = AsyncMock()
|
||
await manager.activate(node_id, websocket)
|
||
await manager.close(node_id)
|
||
websocket.close.assert_awaited_once_with(code=4003, reason="node disabled")
|
||
self.assertFalse(await manager.remove(node_id, websocket))
|
||
|
||
|
||
class SoftwareNodeMigrationTests(unittest.TestCase):
|
||
def test_0030_upgrade_compiles_as_postgresql_ddl(self) -> None:
|
||
statements: list[str] = []
|
||
|
||
def capture(sql, *multiparams, **params):
|
||
statements.append(str(sql.compile(dialect=postgresql.dialect())))
|
||
|
||
engine = create_mock_engine("postgresql+psycopg://", capture)
|
||
operations = Operations(MigrationContext.configure(engine.connect()))
|
||
migration = importlib.import_module(
|
||
"db.migrations.versions.20260812_2000_0030_compute_nodes"
|
||
)
|
||
with patch.object(migration, "op", operations):
|
||
migration.upgrade()
|
||
|
||
rendered = "\n".join(statements)
|
||
self.assertIn("compute_node_enrollments", rendered)
|
||
self.assertIn("compute_nodes", rendered)
|
||
self.assertIn("ix_compute_nodes_status", rendered)
|
||
|
||
def test_0032_upgrade_compiles_as_postgresql_ddl(self) -> None:
|
||
statements: list[str] = []
|
||
|
||
def capture(sql, *multiparams, **params):
|
||
statements.append(str(sql.compile(dialect=postgresql.dialect())))
|
||
|
||
engine = create_mock_engine("postgresql+psycopg://", capture)
|
||
operations = Operations(MigrationContext.configure(engine.connect()))
|
||
migration = importlib.import_module(
|
||
"db.migrations.versions.20260813_1600_0032_software_jobs"
|
||
)
|
||
with patch.object(migration, "op", operations):
|
||
migration.upgrade()
|
||
|
||
rendered = "\n".join(statements)
|
||
self.assertIn("ALTER TABLE compute_node_enrollments RENAME TO software_node_enrollments", rendered)
|
||
self.assertIn("ALTER TABLE compute_nodes RENAME TO software_nodes", rendered)
|
||
self.assertIn("software_jobs", rendered)
|
||
self.assertIn("uq_software_jobs_user_idempotency", rendered)
|
||
self.assertIn("ix_software_jobs_status_created", rendered)
|
||
|
||
def test_0033_adds_software_job_artifact_source(self) -> None:
|
||
statements: list[str] = []
|
||
|
||
def capture(sql, *multiparams, **params):
|
||
statements.append(str(sql.compile(dialect=postgresql.dialect())))
|
||
|
||
engine = create_mock_engine("postgresql+psycopg://", capture)
|
||
operations = Operations(MigrationContext.configure(engine.connect()))
|
||
migration = importlib.import_module(
|
||
"db.migrations.versions.20260814_0900_0033_artifact_software_job"
|
||
)
|
||
with patch.object(migration, "op", operations):
|
||
migration.upgrade()
|
||
|
||
rendered = "\n".join(statements)
|
||
self.assertIn("software_job_id", rendered)
|
||
self.assertIn("ix_artifacts_software_job_id", rendered)
|
||
self.assertIn("jsonb_array_elements", rendered)
|
||
|
||
|
||
class SoftwareJobProtocolTests(unittest.TestCase):
|
||
def test_succeeded_replay_uses_persisted_manifest_across_layout_versions(self) -> None:
|
||
content_digest = "a" * 64
|
||
submitted = [{
|
||
"artifact_id": "plot_spec",
|
||
"filename": "plot-spec.json",
|
||
"media_type": "application/json",
|
||
"size_bytes": 42,
|
||
"sha256": content_digest,
|
||
}]
|
||
persisted = [{
|
||
**submitted[0],
|
||
"source_artifact_id": "plot_spec",
|
||
"artifact_id": str(uuid4()),
|
||
"path": "origin/old-job/plot-spec.json",
|
||
}]
|
||
context = {"status": "succeeded", "artifact_manifest": persisted}
|
||
|
||
self.assertEqual(replay_succeeded_outputs(context, submitted), persisted)
|
||
self.assertTrue(succeeded_output_upload_matches(
|
||
context,
|
||
"plot_spec",
|
||
size_bytes=42,
|
||
digest=content_digest,
|
||
))
|
||
|
||
def test_succeeded_replay_rejects_changed_output(self) -> None:
|
||
context = {
|
||
"status": "succeeded",
|
||
"artifact_manifest": [{
|
||
"source_artifact_id": "figure_png",
|
||
"filename": "figure.png",
|
||
"media_type": "image/png",
|
||
"size_bytes": 10,
|
||
"sha256": "a" * 64,
|
||
"artifact_id": str(uuid4()),
|
||
"path": "origin/job/figure.png",
|
||
}],
|
||
}
|
||
submitted = [{
|
||
"artifact_id": "figure_png",
|
||
"filename": "figure.png",
|
||
"media_type": "image/png",
|
||
"size_bytes": 10,
|
||
"sha256": "b" * 64,
|
||
}]
|
||
|
||
with self.assertRaisesRegex(Exception, "does not match replay"):
|
||
replay_succeeded_outputs(context, submitted)
|
||
with self.assertRaisesRegex(Exception, "does not match replay"):
|
||
succeeded_output_upload_matches(
|
||
context,
|
||
"figure_png",
|
||
size_bytes=10,
|
||
digest="b" * 64,
|
||
)
|
||
|
||
def test_published_output_distinguishes_artifacts_from_metadata(self) -> None:
|
||
job_id = uuid4()
|
||
self.assertTrue(_published_output_is_valid("origin.plot@v2", job_id, {
|
||
"source_artifact_id": "figure_png",
|
||
"artifact_id": str(uuid4()),
|
||
"path": f"origin/{job_id}/figure.png",
|
||
}))
|
||
self.assertTrue(_published_output_is_valid("origin.plot@v2", job_id, {
|
||
"source_artifact_id": "plot_spec",
|
||
"artifact_id": None,
|
||
"path": f"origin/{job_id}/.meta/plot-spec.json",
|
||
}))
|
||
self.assertFalse(_published_output_is_valid("origin.plot@v2", job_id, {
|
||
"source_artifact_id": ["plot_spec"],
|
||
"artifact_id": None,
|
||
"path": f"origin/{job_id}/.meta/plot-spec.json",
|
||
}))
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_queued_job_cancels_without_node_message(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
job = type("Job", (), {})()
|
||
job.job_id = uuid4(); job.user_id = uuid4(); job.task_id = uuid4()
|
||
job.capability = "origin.plot@v2"; job.request_digest = "a" * 64
|
||
job.node_id = None; job.lease_id = None; job.status = "queued"; job.stage = ""
|
||
job.progress = 0; job.metrics = {}; job.error = {}; job.artifact_manifest = []
|
||
job.created_at = None; job.started_at = None; job.terminal_at = None
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
result, node_message = request_job_cancel(job.user_id, job.job_id)
|
||
self.assertEqual(result["status"], "cancelled")
|
||
self.assertIsNone(node_message)
|
||
self.assertEqual(job.error["code"], "USER_CANCELLED")
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_running_job_persists_cancel_before_sending(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
job = type("Job", (), {})()
|
||
job.job_id = uuid4(); job.user_id = uuid4(); job.task_id = uuid4()
|
||
job.capability = "origin.plot@v2"; job.request_digest = "b" * 64
|
||
job.node_id = uuid4(); job.lease_id = uuid4(); job.status = "running"
|
||
job.stage = "software_running"; job.progress = 10; job.metrics = {}; job.error = {}
|
||
job.artifact_manifest = []; job.created_at = None; job.started_at = None; job.terminal_at = None
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
result, node_message = request_job_cancel(job.user_id, job.job_id)
|
||
self.assertEqual(result["status"], "cancelling")
|
||
self.assertEqual(node_message["node_id"], job.node_id)
|
||
self.assertEqual(node_message["payload"]["lease_id"], str(job.lease_id))
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_job_list_is_enriched_for_job_center(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
job = type("Job", (), {})()
|
||
job.job_id = uuid4(); job.task_id = uuid4(); job.capability = "origin.plot@v2"
|
||
job.request_digest = "c" * 64; job.node_id = uuid4(); job.status = "running"
|
||
job.stage = "software_running"; job.progress = 20; job.metrics = {
|
||
"execution_runtime": {
|
||
"node_version": "0.2.0", "adapter_version": "0.8.0",
|
||
"software": "OriginPro", "software_version": "2024",
|
||
}
|
||
}; job.error = {}
|
||
job.artifact_manifest = []; job.input_manifest = [{"key": "sample", "filename": "input.xlsx"}]
|
||
job.request = {"operation": {"plot": {"title": "Test"}}, "outputs": [
|
||
{"key": "figure_png", "type": "figure", "format": "png"}
|
||
]}
|
||
job.created_at = None; job.started_at = None; job.terminal_at = None
|
||
session.execute.return_value.all.return_value = [(job, "材料仿真", "LAB-01")]
|
||
results = list_jobs(uuid4(), limit=10)
|
||
self.assertEqual(results[0]["task_name"], "材料仿真")
|
||
self.assertEqual(results[0]["node_name"], "LAB-01")
|
||
self.assertEqual(results[0]["request_summary"]["display_name"], "Origin 科研绘图")
|
||
self.assertEqual(results[0]["request_summary"]["formats"], ["png"])
|
||
self.assertEqual(results[0]["output_dir"], f"origin/{job.job_id}")
|
||
self.assertEqual(results[0]["execution_runtime"]["software_version"], "2024")
|
||
|
||
def test_origin_request_is_canonical_and_rejects_extra_fields(self) -> None:
|
||
request = {
|
||
"schema_version": 2,
|
||
"inputs": [
|
||
{"key": "first", "artifact_id": str(uuid4()), "selector": {"sheet": "Sheet1"}},
|
||
{"key": "second", "artifact_id": str(uuid4())},
|
||
],
|
||
"operation": {"plot": {
|
||
"type": "line",
|
||
"series": [
|
||
{"input": "first", "x": "x", "y": "y"},
|
||
{"input": "second", "x": "time", "y": "value", "label": "Second"},
|
||
],
|
||
}},
|
||
"outputs": [
|
||
{"key": "figure_png", "type": "figure", "format": "png", "options": {"dpi": 600}},
|
||
{"key": "project", "type": "project", "format": "opju"},
|
||
],
|
||
}
|
||
normalized, digest = _canonical_request("origin.plot@v2", request)
|
||
self.assertEqual(normalized, request)
|
||
self.assertEqual(len(digest), 64)
|
||
with self.assertRaisesRegex(Exception, "invalid origin.plot@v2 request"):
|
||
_canonical_request("origin.plot@v2", {**request, "script": "anything"})
|
||
with self.assertRaisesRegex(Exception, "invalid origin.plot@v2 request"):
|
||
_canonical_request("origin.plot@v2", {**request, "operation": {"plot": {
|
||
**request["operation"]["plot"], "script": "anything"
|
||
}}})
|
||
with self.assertRaisesRegex(Exception, "not a 'uuid'"):
|
||
_canonical_request("origin.plot@v2", {**request, "inputs": [{"key": "first", "artifact_id": "C:\\data.csv"}]})
|
||
|
||
def test_origin_request_rejects_unimplemented_plot_semantics(self) -> None:
|
||
request = {
|
||
"schema_version": 2,
|
||
"inputs": [{"key": "sample", "artifact_id": str(uuid4())}],
|
||
"operation": {"plot": {"type": "scatter", "series": [
|
||
{"input": "sample", "x": "time", "y": "a"},
|
||
{"input": "sample", "x": "time", "y": "b"},
|
||
]}},
|
||
"outputs": [{"key": "figure_png", "type": "figure", "format": "png", "options": {"dpi": 300}}],
|
||
}
|
||
base_plot = request["operation"]["plot"]
|
||
for case_plot in (
|
||
{**base_plot, "template": "custom"},
|
||
{**base_plot, "x_axis": {"scale": "symlog"}},
|
||
{**base_plot, "legend": {"position": "outside"}},
|
||
{**base_plot, "series": [base_plot["series"][0], base_plot["series"][0]]},
|
||
):
|
||
with self.subTest(plot=case_plot), self.assertRaisesRegex(
|
||
Exception, "invalid origin.plot@v2 request"
|
||
):
|
||
_canonical_request("origin.plot@v2", {**request, "operation": {"plot": case_plot}})
|
||
# 跨字段引用完整性属于本机 adapter 的独立二次校验；Core 只执行共享结构契约。
|
||
normalized, _ = _canonical_request("origin.plot@v2", {
|
||
**request,
|
||
"operation": {"plot": {
|
||
**base_plot,
|
||
"series": [{"input": "missing", "x": "time", "y": "a"}],
|
||
}},
|
||
})
|
||
self.assertEqual(normalized["operation"]["plot"]["series"][0]["input"], "missing")
|
||
with self.assertRaisesRegex(Exception, "invalid origin.plot@v2 request"):
|
||
_canonical_request(
|
||
"origin.plot@v2",
|
||
{**request, "outputs": [{"key": "project", "type": "figure", "format": "png"}]}
|
||
)
|
||
|
||
def test_origin_request_accepts_publication_layout_options(self) -> None:
|
||
request = {
|
||
"schema_version": 2,
|
||
"inputs": [{"key": "sample", "artifact_id": str(uuid4())}],
|
||
"operation": {"plot": {
|
||
"type": "line_scatter",
|
||
"series": [{
|
||
"input": "sample", "x": "time", "y": "strength",
|
||
"style": {
|
||
"color": "#3366CC", "line_width": 1.5,
|
||
"line_style": "dash", "symbol": "circle",
|
||
"symbol_size": 8, "transparency": 10,
|
||
},
|
||
}],
|
||
"title": "Strength development",
|
||
"title_style": {"font_size": 16},
|
||
"canvas": {"width_mm": 180, "height_mm": 120},
|
||
"x_axis": {
|
||
"scale": "log10", "minimum": 1, "maximum": 100,
|
||
"major_step": 1, "tick_label_angle": 45,
|
||
"tick_label_font_size": 10, "title_font_size": 12,
|
||
"grid": "major",
|
||
},
|
||
"y_axis": {"minimum": 0, "maximum": 80, "major_step": 10},
|
||
"legend": {"enabled": False, "position": "auto", "font_size": 9},
|
||
}},
|
||
"outputs": [{
|
||
"key": "figure_png", "type": "figure", "format": "png",
|
||
"options": {"dpi": 600},
|
||
}],
|
||
}
|
||
normalized, digest = _canonical_request("origin.plot@v2", request)
|
||
self.assertEqual(normalized, request)
|
||
self.assertEqual(len(digest), 64)
|
||
|
||
def test_origin_request_supports_unified_series_roles(self) -> None:
|
||
artifact_id = str(uuid4())
|
||
outputs = [{"key": "figure_png", "type": "figure", "format": "png"}]
|
||
cases = {
|
||
"column": [{"input": "sample", "x": "age", "y": "strength"}],
|
||
"bar": [{"input": "sample", "x": "age", "y": "strength"}],
|
||
"grouped_column": [
|
||
{"input": "sample", "x": "age", "y": "strength"},
|
||
{"input": "sample", "x": "age", "y": "modulus"},
|
||
],
|
||
"y_error": [{"input": "sample", "x": "age", "y": "strength", "y_error": "sd"}],
|
||
"contour": [{"input": "sample", "x": "x", "y": "y", "z": "value"}],
|
||
"surface_3d": [{"input": "sample", "x": "x", "y": "y", "z": "value"}],
|
||
"ternary": [{"input": "sample", "x": "a", "y": "b", "z": "c"}],
|
||
"heatmap": [{"input": "sample", "x": "x", "y": "y", "z": "value"}],
|
||
"box": [{"input": "sample", "y": "strength"}],
|
||
"histogram": [{"input": "sample", "y": "particle_size"}],
|
||
"stacked_column": [
|
||
{"input": "sample", "x": "sample_name", "y": "phase_a"},
|
||
{"input": "sample", "x": "sample_name", "y": "phase_b"},
|
||
],
|
||
"bubble": [{
|
||
"input": "sample", "x": "temperature", "y": "strength",
|
||
"size": "particle_size",
|
||
}],
|
||
"band": [{
|
||
"input": "sample", "x": "age", "y": "mean",
|
||
"lower": "ci_lower", "upper": "ci_upper",
|
||
}],
|
||
"area": [{"input": "sample", "x": "age", "y": "strength"}],
|
||
"stacked_area": [
|
||
{"input": "sample", "x": "age", "y": "phase_a"},
|
||
{"input": "sample", "x": "age", "y": "phase_b"},
|
||
],
|
||
"polar": [{"input": "sample", "x": "angle", "y": "intensity"}],
|
||
"pie": [{"input": "sample", "x": "phase", "y": "fraction"}],
|
||
"stacked_bar": [
|
||
{"input": "sample", "x": "sample_name", "y": "phase_a"},
|
||
{"input": "sample", "x": "sample_name", "y": "phase_b"},
|
||
],
|
||
}
|
||
for plot_type, series in cases.items():
|
||
request = {
|
||
"schema_version": 2,
|
||
"inputs": [{"key": "sample", "artifact_id": artifact_id}],
|
||
"operation": {"plot": {"type": plot_type, "series": series}},
|
||
"outputs": outputs,
|
||
}
|
||
with self.subTest(plot_type=plot_type):
|
||
normalized, digest = _canonical_request("origin.plot@v2", request)
|
||
self.assertEqual(normalized, request)
|
||
self.assertEqual(len(digest), 64)
|
||
|
||
def test_origin_request_accepts_multi_panel_mixed_series_and_dual_y(self) -> None:
|
||
request = {
|
||
"schema_version": 2,
|
||
"inputs": [{"key": "sample", "artifact_id": str(uuid4())}],
|
||
"operation": {"plot": {
|
||
"type": "multi_panel",
|
||
"layout": {"rows": 1, "columns": 2, "share_x": False},
|
||
"canvas": {"width_mm": 180, "height_mm": 90},
|
||
"panels": [
|
||
{
|
||
"key": "mechanical",
|
||
"panel_label": "(a)",
|
||
"title": "Mechanical properties",
|
||
"series": [
|
||
{
|
||
"input": "sample", "x": "age", "y": "strength",
|
||
"kind": "column", "y_error": "strength_sd",
|
||
},
|
||
{
|
||
"input": "sample", "x": "age", "y": "modulus",
|
||
"kind": "line_scatter", "y_axis": "right",
|
||
},
|
||
],
|
||
"x_axis": {"title": "Age", "unit": "d"},
|
||
"y_axis": {"title": "Strength", "unit": "MPa"},
|
||
"right_y_axis": {"title": "Modulus", "unit": "GPa"},
|
||
},
|
||
{
|
||
"key": "fit",
|
||
"panel_label": "(b)",
|
||
"series": [
|
||
{
|
||
"input": "sample", "x": "measured", "y": "predicted",
|
||
"kind": "scatter", "x_error": "measured_sd",
|
||
},
|
||
{
|
||
"input": "sample", "x": "measured", "y": "reference",
|
||
"kind": "line",
|
||
},
|
||
],
|
||
},
|
||
],
|
||
}},
|
||
"outputs": [{"key": "figure_svg", "type": "figure", "format": "svg"}],
|
||
}
|
||
normalized, digest = _canonical_request("origin.plot@v2", request)
|
||
self.assertEqual(normalized, request)
|
||
self.assertEqual(len(digest), 64)
|
||
|
||
def test_origin_request_rejects_roles_that_do_not_match_plot_type(self) -> None:
|
||
base = {
|
||
"schema_version": 2,
|
||
"inputs": [{"key": "sample", "artifact_id": str(uuid4())}],
|
||
"outputs": [{"key": "figure_png", "type": "figure", "format": "png"}],
|
||
}
|
||
invalid_plots = (
|
||
{"type": "line", "series": [{"input": "sample", "x": "x", "y": "y", "z": "z"}]},
|
||
{"type": "line", "series": [{
|
||
"input": "sample", "x": "x", "y": "y", "size": "size",
|
||
}]},
|
||
{"type": "y_error", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "contour", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "grouped_column", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "box", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "stacked_column", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "bubble", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "band", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "stacked_area", "series": [{"input": "sample", "x": "x", "y": "y"}]},
|
||
{"type": "polar", "series": [{"input": "sample", "x": "x", "y": "y"}], "x_axis": {"title": "Angle"}},
|
||
)
|
||
for plot in invalid_plots:
|
||
with self.subTest(plot_type=plot["type"]), self.assertRaises(SoftwareJobError):
|
||
_canonical_request("origin.plot@v2", {**base, "operation": {"plot": plot}})
|
||
|
||
def test_output_manifest_matches_exact_requested_outputs(self) -> None:
|
||
request = {"outputs": [
|
||
{"key": "project", "type": "project", "format": "opju"},
|
||
{"key": "figure_png", "type": "figure", "format": "png"},
|
||
]}
|
||
manifest = [
|
||
{"artifact_id": "project", "filename": "project.opju", "media_type": "application/x-origin-project", "size_bytes": 10, "sha256": "a" * 64},
|
||
{"artifact_id": "figure_png", "filename": "figure.png", "media_type": "image/png", "size_bytes": 20, "sha256": "b" * 64},
|
||
{"artifact_id": "plot_spec", "filename": "plot-spec.json", "media_type": "application/json", "size_bytes": 30, "sha256": "c" * 64},
|
||
{"artifact_id": "provenance", "filename": "provenance.json", "media_type": "application/json", "size_bytes": 40, "sha256": "d" * 64},
|
||
]
|
||
self.assertEqual(
|
||
validate_output_manifest("origin.plot@v2", request, manifest), manifest
|
||
)
|
||
with self.assertRaisesRegex(Exception, "incomplete"):
|
||
validate_output_manifest("origin.plot@v2", request, manifest[:-1])
|
||
with self.assertRaisesRegex(Exception, "metadata"):
|
||
validate_output_manifest(
|
||
"origin.plot@v2", request,
|
||
[{**manifest[0], "filename": "anything.opju"}, *manifest[1:]],
|
||
)
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_stale_offer_cannot_be_accepted_by_another_node(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
job = type("Job", (), {})()
|
||
job.node_id = uuid4()
|
||
job.lease_id = uuid4()
|
||
job.status = "offered"
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
with self.assertRaisesRegex(Exception, "stale or does not belong"):
|
||
respond_to_offer(
|
||
uuid4(),
|
||
accepted=True,
|
||
payload={"job_id": str(uuid4()), "lease_id": str(job.lease_id)},
|
||
)
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_accepted_job_snapshots_execution_versions(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
node_id = uuid4()
|
||
lease_id = uuid4()
|
||
digest = "d" * 64
|
||
job = type("Job", (), {})()
|
||
job.node_id = node_id; job.lease_id = lease_id; job.status = "offered"
|
||
job.lease_expires_at = datetime.now(timezone.utc) + timedelta(minutes=1)
|
||
job.request_digest = digest; job.capability = "origin.plot@v2"; job.metrics = {}
|
||
node = type("Node", (), {})()
|
||
node.node_version = "0.1.0"
|
||
node.runtime = {
|
||
"node_version": "0.2.0",
|
||
"capability_runtime": {"origin.plot@v2": {
|
||
"adapter_version": "0.8.0", "health": "ready",
|
||
}},
|
||
"origin": {
|
||
"software": "OriginPro", "software_version": "2024",
|
||
"adapter_version": "0.8.0", "health": "ready",
|
||
},
|
||
}
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
session.get.return_value = node
|
||
|
||
respond_to_offer(node_id, accepted=True, payload={
|
||
"job_id": str(uuid4()), "lease_id": str(lease_id),
|
||
"request_digest": digest,
|
||
})
|
||
|
||
self.assertEqual(job.metrics["execution_runtime"], {
|
||
"node_version": "0.2.0", "adapter_version": "0.8.0",
|
||
"software": "OriginPro", "software_version": "2024",
|
||
})
|
||
|
||
job.started_at = None
|
||
update_job_state(node_id, {
|
||
"job_id": str(uuid4()), "lease_id": str(lease_id),
|
||
"request_digest": digest, "stage": "software_running",
|
||
"progress": 10, "metrics": {"elapsed_seconds": 1},
|
||
})
|
||
self.assertEqual(job.metrics["elapsed_seconds"], 1)
|
||
self.assertEqual(job.metrics["execution_runtime"]["software_version"], "2024")
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_failed_delivery_only_abandons_matching_offer(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
node_id = uuid4()
|
||
lease_id = uuid4()
|
||
job = type("Job", (), {})()
|
||
job.node_id = node_id
|
||
job.lease_id = lease_id
|
||
job.status = "offered"
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
abandon_offer(
|
||
node_id,
|
||
{"job_id": str(uuid4()), "lease_id": str(lease_id)},
|
||
)
|
||
self.assertEqual(job.status, "queued")
|
||
self.assertIsNone(job.node_id)
|
||
|
||
def test_dispatcher_excludes_nodes_with_active_jobs(self) -> None:
|
||
source = (
|
||
Path(__file__).resolve().parents[1] / "core" / "software_jobs.py"
|
||
).read_text(encoding="utf-8")
|
||
self.assertIn('"disconnected", "cancelling"', source)
|
||
self.assertIn("node_supports_request", source)
|
||
self.assertIn("SoftwareJob.capability.in_(available_capabilities)", source)
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_dispatcher_skips_job_that_requires_newer_adapter(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
node = type("Node", (), {})()
|
||
node.node_id = uuid4()
|
||
node.capabilities = ["origin.plot@v2"]
|
||
node.runtime = {"capability_runtime": {"origin.plot@v2": {
|
||
"health": "ready", "available_slots": 1, "adapter_version": "0.5.0",
|
||
"features": ["line"],
|
||
}}}
|
||
jobs = []
|
||
for plot_type in ("heatmap", "line"):
|
||
job = type("Job", (), {})()
|
||
job.job_id = uuid4()
|
||
job.capability = "origin.plot@v2"
|
||
roles = {"input": "sample", "x": "x", "y": "y"}
|
||
if plot_type == "heatmap":
|
||
roles["z"] = "z"
|
||
job.request = {
|
||
"operation": {"plot": {"type": plot_type, "series": [roles]}},
|
||
"outputs": [],
|
||
}
|
||
job.input_manifest = []
|
||
job.request_digest = plot_type[0] * 64
|
||
job.status = "queued"
|
||
jobs.append(job)
|
||
results = [MagicMock() for _ in range(5)]
|
||
results[0].scalars.return_value = []
|
||
results[1].scalars.return_value = []
|
||
results[2].scalars.return_value = [node]
|
||
results[3].scalars.return_value = jobs
|
||
results[4].scalar_one_or_none.return_value = jobs[1]
|
||
session.execute.side_effect = results
|
||
|
||
offer = offer_next_job({node.node_id})
|
||
|
||
self.assertEqual(offer["node_id"], node.node_id)
|
||
self.assertEqual(offer["payload"]["job_id"], str(jobs[1].job_id))
|
||
self.assertEqual(offer["payload"]["request"]["operation"]["plot"]["type"], "line")
|
||
|
||
def test_input_download_rechecks_file_digest(self) -> None:
|
||
source = (
|
||
Path(__file__).resolve().parents[1]
|
||
/ "web" / "routers" / "software_nodes.py"
|
||
).read_text(encoding="utf-8")
|
||
self.assertIn("digest = sha256()", source)
|
||
self.assertIn('digest.hexdigest() != item["sha256"]', source)
|
||
self.assertIn('contract.expected_outputs(context["request"])', source)
|
||
self.assertIn("artifact_id not in requested_ids", source)
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_job_state_restores_disconnected_job(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
node_id = uuid4()
|
||
lease_id = uuid4()
|
||
digest = "a" * 64
|
||
job = type("Job", (), {})()
|
||
job.node_id = node_id
|
||
job.lease_id = lease_id
|
||
job.request_digest = digest
|
||
job.status = "disconnected"
|
||
job.started_at = None
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
update_job_state(node_id, {
|
||
"job_id": str(uuid4()),
|
||
"lease_id": str(lease_id),
|
||
"request_digest": digest,
|
||
"stage": "downloading_inputs",
|
||
"progress": 0,
|
||
"metrics": {},
|
||
})
|
||
self.assertEqual(job.status, "dispatched")
|
||
self.assertEqual(job.stage, "downloading_inputs")
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_ready_to_run_is_not_reported_as_running(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
node_id = uuid4()
|
||
lease_id = uuid4()
|
||
digest = "c" * 64
|
||
job = type("Job", (), {})()
|
||
job.node_id = node_id
|
||
job.lease_id = lease_id
|
||
job.request_digest = digest
|
||
job.status = "dispatched"
|
||
job.started_at = None
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
update_job_state(node_id, {
|
||
"job_id": str(uuid4()), "lease_id": str(lease_id),
|
||
"request_digest": digest, "stage": "ready_to_run",
|
||
"progress": 5, "metrics": {"input_bytes": 10},
|
||
})
|
||
self.assertEqual(job.status, "dispatched")
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_terminal_replay_is_idempotent(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
node_id = uuid4()
|
||
lease_id = uuid4()
|
||
digest = "b" * 64
|
||
job = type("Job", (), {})()
|
||
job.node_id = node_id
|
||
job.lease_id = lease_id
|
||
job.request_digest = digest
|
||
job.status = "failed"
|
||
session.execute.return_value.scalar_one_or_none.return_value = job
|
||
record_job_terminal(node_id, {
|
||
"job_id": str(uuid4()),
|
||
"lease_id": str(lease_id),
|
||
"request_digest": digest,
|
||
"status": "failed",
|
||
"error": {"code": "TEST"},
|
||
"artifact_manifest": [],
|
||
})
|
||
self.assertEqual(job.status, "failed")
|
||
|
||
@patch("core.software_jobs.session_scope")
|
||
def test_disconnect_does_not_requeue_active_jobs(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
first = type("Job", (), {"status": "running"})()
|
||
second = type("Job", (), {"status": "dispatched"})()
|
||
session.execute.return_value.scalars.return_value = [first, second]
|
||
mark_node_jobs_disconnected(uuid4())
|
||
self.assertEqual(first.status, "disconnected")
|
||
self.assertEqual(second.status, "disconnected")
|
||
|
||
|
||
class SoftwareNodeDeleteTests(unittest.TestCase):
|
||
@patch("core.software_nodes.session_scope")
|
||
def test_delete_node_removes_existing_identity(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
node = object()
|
||
session.get.return_value = node
|
||
|
||
self.assertTrue(delete_node(uuid4()))
|
||
session.delete.assert_called_once_with(node)
|
||
|
||
@patch("core.software_nodes.session_scope")
|
||
def test_delete_node_reports_missing_identity(self, session_scope) -> None:
|
||
session = session_scope.return_value.__enter__.return_value
|
||
session.get.return_value = None
|
||
|
||
self.assertFalse(delete_node(uuid4()))
|
||
session.delete.assert_not_called()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|