310 lines
12 KiB
Python
310 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from k1link.simulation.gaussian_pipeline_gateway import (
|
|
BUILD_REQUEST_SCHEMA,
|
|
GaussianPipelineGateway,
|
|
GaussianPipelineGatewayError,
|
|
GaussianPipelineIntegrityError,
|
|
_discover_bundle_members,
|
|
)
|
|
|
|
|
|
def _token_file(tmp_path: Path) -> Path:
|
|
token = tmp_path / "gaussian.token"
|
|
token.write_text("t" * 64, encoding="utf-8")
|
|
return token
|
|
|
|
|
|
def test_gateway_uploads_lcc_bundle_with_tus_and_reads_provider_contract(tmp_path: Path) -> None:
|
|
uploads: dict[str, bytearray] = {}
|
|
metadata: dict[str, dict[str, str]] = {}
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
assert request.headers["authorization"] == f"Bearer {'t' * 64}"
|
|
if request.method == "GET" and request.url.path == "/v1/capabilities":
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"schema_version": "gaussian-pipeline.capabilities/v1",
|
|
"service": "ndc-gaussian-pipeline",
|
|
"api_version": "gaussian-pipeline.api/v1",
|
|
"upload_protocol": "tus/1.0.0",
|
|
"source_transport": "tus-bundle/v1",
|
|
"provider": {"id": "playcanvas-splat-transform", "version": "3.3.3"},
|
|
"runtime": {
|
|
"source_revision": "a" * 40,
|
|
"image_digest": f"sha256:{'b' * 64}",
|
|
},
|
|
"max_source_bytes": 1024 * 1024,
|
|
"max_source_files": 10_000,
|
|
},
|
|
)
|
|
if request.method == "POST" and request.url.path == "/v1/uploads":
|
|
upload_id = f"upload-{len(uploads) + 1:03d}"
|
|
uploads[upload_id] = bytearray()
|
|
item_metadata: dict[str, str] = {}
|
|
for item in request.headers["upload-metadata"].split(","):
|
|
key, encoded = item.split(" ", 1)
|
|
item_metadata[key] = base64.b64decode(encoded).decode("utf-8")
|
|
metadata[upload_id] = item_metadata
|
|
return httpx.Response(201, headers={"Location": f"/v1/uploads/{upload_id}"})
|
|
upload_id = request.url.path.rsplit("/", 1)[-1]
|
|
if request.method == "HEAD" and upload_id in uploads:
|
|
return httpx.Response(200, headers={"Upload-Offset": str(len(uploads[upload_id]))})
|
|
if request.method == "PATCH" and upload_id in uploads:
|
|
assert int(request.headers["upload-offset"]) == len(uploads[upload_id])
|
|
uploads[upload_id].extend(request.content)
|
|
return httpx.Response(
|
|
204,
|
|
headers={"Upload-Offset": str(len(uploads[upload_id]))},
|
|
)
|
|
return httpx.Response(404)
|
|
|
|
bundle = tmp_path / "bundle"
|
|
bundle.mkdir()
|
|
(bundle / "yard.lcc").write_text(
|
|
json.dumps({"fileType": "Quality", "attributes": []}),
|
|
encoding="utf-8",
|
|
)
|
|
(bundle / "index.bin").write_bytes(b"index")
|
|
(bundle / "data.bin").write_bytes(b"data")
|
|
(bundle / "shcoef.bin").write_bytes(b"sh")
|
|
with GaussianPipelineGateway(
|
|
"http://gaussian.test",
|
|
_token_file(tmp_path),
|
|
chunk_bytes=5,
|
|
transport=httpx.MockTransport(handler),
|
|
) as gateway:
|
|
assert gateway.capabilities()["service"] == "ndc-gaussian-pipeline"
|
|
descriptor = gateway.upload_source_bundle(
|
|
bundle,
|
|
entrypoint="yard.lcc",
|
|
source_format="lcc",
|
|
)
|
|
|
|
paths = [member.logical_path for member in descriptor.members]
|
|
assert paths == ["data.bin", "index.bin", "shcoef.bin", "yard.lcc"]
|
|
assert [metadata[f"upload-{index:03d}"]["logical_path"] for index in range(1, 5)] == paths
|
|
for index, member in enumerate(descriptor.members, start=1):
|
|
assert bytes(uploads[f"upload-{index:03d}"]) == (bundle / member.logical_path).read_bytes()
|
|
assert metadata[f"upload-{index:03d}"]["sha256"] == member.sha256
|
|
identity = {
|
|
"format": "lcc",
|
|
"entrypoint": "yard.lcc",
|
|
"members": [
|
|
{
|
|
"logical_path": member.logical_path,
|
|
"sha256": member.sha256,
|
|
"byte_length": member.byte_length,
|
|
}
|
|
for member in descriptor.members
|
|
],
|
|
}
|
|
assert descriptor.bundle_sha256 == hashlib.sha256(
|
|
json.dumps(identity, ensure_ascii=False, separators=(",", ":")).encode()
|
|
).hexdigest()
|
|
assert descriptor.total_byte_length == sum(member.byte_length for member in descriptor.members)
|
|
|
|
|
|
def test_gateway_rejects_incomplete_lcc_bundle(tmp_path: Path) -> None:
|
|
bundle = tmp_path / "bundle"
|
|
bundle.mkdir()
|
|
(bundle / "yard.lcc").write_text('{"fileType":"Portable"}', encoding="utf-8")
|
|
with (
|
|
GaussianPipelineGateway(
|
|
"http://gaussian.test",
|
|
_token_file(tmp_path),
|
|
transport=httpx.MockTransport(lambda _request: httpx.Response(500)),
|
|
) as gateway,
|
|
pytest.raises(GaussianPipelineIntegrityError, match="member is unavailable"),
|
|
):
|
|
gateway.upload_source_bundle(
|
|
bundle,
|
|
entrypoint="yard.lcc",
|
|
source_format="lcc",
|
|
)
|
|
|
|
|
|
def test_gateway_discovers_nested_lcc2_chunks_with_trailing_commas(tmp_path: Path) -> None:
|
|
bundle = tmp_path / "bundle"
|
|
(bundle / "scene" / "chunks").mkdir(parents=True)
|
|
(bundle / "scene" / "meta.lcc2").write_text(
|
|
"""
|
|
{
|
|
"root": {"splatFiles": ["chunks/0.sog", "chunks/1.spz",],},
|
|
}
|
|
""",
|
|
encoding="utf-8",
|
|
)
|
|
(bundle / "scene" / "chunks" / "0.sog").write_bytes(b"sog")
|
|
(bundle / "scene" / "chunks" / "1.spz").write_bytes(b"spz")
|
|
|
|
assert _discover_bundle_members(bundle, "scene/meta.lcc2", "lcc2") == {
|
|
"scene/meta.lcc2",
|
|
"scene/chunks/0.sog",
|
|
"scene/chunks/1.spz",
|
|
}
|
|
|
|
|
|
def test_gateway_rejects_upload_location_outside_provider(tmp_path: Path) -> None:
|
|
bundle = tmp_path / "bundle"
|
|
bundle.mkdir()
|
|
(bundle / "yard.lcc").write_text('{"fileType":"Portable"}', encoding="utf-8")
|
|
(bundle / "data.bin").write_bytes(b"data")
|
|
(bundle / "index.bin").write_bytes(b"index")
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
if request.method == "GET":
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"schema_version": "gaussian-pipeline.capabilities/v1",
|
|
"service": "ndc-gaussian-pipeline",
|
|
"api_version": "gaussian-pipeline.api/v1",
|
|
"upload_protocol": "tus/1.0.0",
|
|
"source_transport": "tus-bundle/v1",
|
|
"runtime": {
|
|
"source_revision": "a" * 40,
|
|
"image_digest": f"sha256:{'b' * 64}",
|
|
},
|
|
"max_source_bytes": 1024,
|
|
"max_source_files": 10_000,
|
|
},
|
|
)
|
|
return httpx.Response(
|
|
201,
|
|
headers={"Location": "https://attacker.invalid/v1/uploads/stolen"},
|
|
)
|
|
|
|
with (
|
|
GaussianPipelineGateway(
|
|
"http://gaussian.test",
|
|
_token_file(tmp_path),
|
|
transport=httpx.MockTransport(handler),
|
|
) as gateway,
|
|
pytest.raises(GaussianPipelineGatewayError, match="escaped the provider"),
|
|
):
|
|
gateway.upload_source_bundle(
|
|
bundle,
|
|
entrypoint="yard.lcc",
|
|
source_format="lcc",
|
|
)
|
|
|
|
|
|
def test_gateway_rejects_symlink_source_and_token(tmp_path: Path) -> None:
|
|
bundle = tmp_path / "bundle"
|
|
bundle.mkdir()
|
|
(bundle / "source.lcc").write_text('{"fileType":"Portable"}', encoding="utf-8")
|
|
(bundle / "real-data.bin").write_bytes(b"source")
|
|
(bundle / "data.bin").symlink_to(bundle / "real-data.bin")
|
|
(bundle / "index.bin").write_bytes(b"index")
|
|
token = _token_file(tmp_path)
|
|
token_link = tmp_path / "token-link"
|
|
token_link.symlink_to(token)
|
|
|
|
with pytest.raises(GaussianPipelineGatewayError, match="token must be one regular file"):
|
|
GaussianPipelineGateway("http://gaussian.test", token_link)
|
|
|
|
with (
|
|
GaussianPipelineGateway(
|
|
"http://gaussian.test",
|
|
token,
|
|
transport=httpx.MockTransport(lambda _request: httpx.Response(500)),
|
|
) as gateway,
|
|
pytest.raises(GaussianPipelineIntegrityError, match="contains a symlink"),
|
|
):
|
|
gateway.upload_source_bundle(
|
|
bundle,
|
|
entrypoint="source.lcc",
|
|
source_format="lcc",
|
|
)
|
|
|
|
|
|
def test_gateway_submits_tracks_and_imports_digest_bound_artifact(tmp_path: Path) -> None:
|
|
artifact = b"preview-sog"
|
|
artifact_sha = hashlib.sha256(artifact).hexdigest()
|
|
job_id = "gsp-20260825200000-deadbeef"
|
|
request_document: dict[str, object] = {
|
|
"schema_version": BUILD_REQUEST_SCHEMA,
|
|
"idempotency_key": "missioncore-build-01",
|
|
}
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
if request.method == "POST" and request.url.path == "/v1/jobs":
|
|
assert json.loads(request.content) == request_document
|
|
return httpx.Response(
|
|
202,
|
|
json={"schema_version": "gaussian-pipeline.job/v1", "job_id": job_id},
|
|
)
|
|
if request.method == "GET" and request.url.path == f"/v1/jobs/{job_id}":
|
|
return httpx.Response(
|
|
200,
|
|
json={"schema_version": "gaussian-pipeline.job/v1", "job_id": job_id},
|
|
)
|
|
if request.method == "GET" and request.url.path == f"/v1/jobs/{job_id}/result":
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"schema_version": "gaussian-pipeline.build-result/v1",
|
|
"job_id": job_id,
|
|
"runtime": {
|
|
"source_revision": "a" * 40,
|
|
"image_digest": f"sha256:{'b' * 64}",
|
|
},
|
|
},
|
|
)
|
|
if request.method == "GET" and request.url.path.endswith("/artifacts/preview.sog"):
|
|
return httpx.Response(200, content=artifact)
|
|
return httpx.Response(404)
|
|
|
|
with GaussianPipelineGateway(
|
|
"http://gaussian.test",
|
|
_token_file(tmp_path),
|
|
transport=httpx.MockTransport(handler),
|
|
) as gateway:
|
|
assert gateway.submit_build(request_document)["job_id"] == job_id
|
|
assert gateway.get_job(job_id)["job_id"] == job_id
|
|
assert gateway.get_result(job_id)["job_id"] == job_id
|
|
destination = gateway.download_artifact(
|
|
job_id,
|
|
{
|
|
"logical_path": "preview.sog",
|
|
"sha256": artifact_sha,
|
|
"byte_length": len(artifact),
|
|
},
|
|
tmp_path / "import" / "preview.sog",
|
|
)
|
|
assert destination.read_bytes() == artifact
|
|
|
|
|
|
def test_gateway_rejects_capabilities_without_runtime_provenance(tmp_path: Path) -> None:
|
|
def handler(_request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
200,
|
|
json={
|
|
"schema_version": "gaussian-pipeline.capabilities/v1",
|
|
"service": "ndc-gaussian-pipeline",
|
|
"api_version": "gaussian-pipeline.api/v1",
|
|
"upload_protocol": "tus/1.0.0",
|
|
"source_transport": "tus-bundle/v1",
|
|
},
|
|
)
|
|
|
|
with (
|
|
GaussianPipelineGateway(
|
|
"http://gaussian.test",
|
|
_token_file(tmp_path),
|
|
transport=httpx.MockTransport(handler),
|
|
) as gateway,
|
|
pytest.raises(GaussianPipelineGatewayError, match="provenance is unavailable"),
|
|
):
|
|
gateway.capabilities()
|