story: DGR-041 Register native Shard capabilities without redesigning Meshnet

tests/test_native_registration.py

@@ -0,0 +1,94 @@
+"""DGR-041 native capability registration remains an ordinary admission payload."""
+
+from __future__ import annotations
+
+from meshnet_node.capability import ExecutionCapacity
+from meshnet_node.native_registration import (
+    NativeCapabilityRegistrar,
+    NativeRegistrationError,
+    NativeShardRegistration,
+)
+from meshnet_node.native_worker_supervisor import NativeWorkerProbe, NativeWorkerSpec
+from meshnet_node.runtime_recipe import ShardIdentity
+from meshnet_tracker.capability import STATE_UNCERTIFIED
+from meshnet_tracker.server import TrackerServer, _capability_from_registration
+
+from test_runtime_recipe_identity import _identity
+
+
+def _worker(identity: ShardIdentity) -> tuple[NativeWorkerSpec, NativeWorkerProbe]:
+    spec = NativeWorkerSpec(
+        binary=__file__, binary_digest="d" * 64, listen_address="127.0.0.1:1",
+        artifact_path=__file__, artifact_digest=identity.fingerprint.model_artifact_digest,
+        recipe_digest=identity.fingerprint.runtime_recipe_digest, recipe_id=identity.recipe.recipe_id,
+        recipe_version=identity.recipe.recipe_version, catalogue_version=identity.recipe.catalogue_version,
+        shard_start=identity.shard_start, shard_end=identity.shard_end,
+    )
+    probe = NativeWorkerProbe(
+        artifact_digest=spec.artifact_digest, recipe_digest=spec.recipe_digest,
+        recipe_id=spec.recipe_id, recipe_version=spec.recipe_version,
+        catalogue_version=spec.catalogue_version, shard_start=spec.shard_start,
+        shard_end=spec.shard_end, serving=True,
+    )
+    return spec, probe
+
+
+def test_native_registration_carries_exact_identity_range_capacity_and_dark_status():
+    identity = _identity()
+    worker, probe = _worker(identity)
+    registration = NativeShardRegistration(
+        endpoint="http://native.example:8080", model_id=identity.artifact.artifact_id,
+        identity=identity, worker=worker, probe=probe, device="cpu:fixture",
+        capacity=ExecutionCapacity(4096, 8192, 3), duration_ms=7,
+    )
+
+    payload = registration.payload()
+    report = payload["capability_report"]
+    assert report["identity"]["fingerprint"]["runtime_recipe_digest"] == identity.fingerprint.runtime_recipe_digest
+    assert report["shard"] == {"start": identity.shard_start, "end": identity.shard_end - 1}
+    assert report["backend"]["backend_id"] == identity.recipe.axes["backend_id"]
+    assert report["capacity"] == {
+        "memory_capacity_bytes": 4096, "kv_capacity_tokens": 8192, "max_concurrent_sessions": 3,
+    }
+
+    tracker = TrackerServer()
+    state = _capability_from_registration(
+        payload, model=payload["model"], hf_repo=payload["hf_repo"],
+        shard_start=payload["shard_start"], shard_end=payload["shard_end"],
+        recipe_certifications=tracker._recipe_certifications,
+    )
+    assert state.state == STATE_UNCERTIFIED
+    assert state.certification == "dark"
+    assert state.memory_capacity_bytes == 4096
+    assert state.kv_capacity_tokens == 8192
+    assert state.max_concurrent_sessions == 3
+
+
+def test_native_registrar_has_no_tracker_or_backend_policy_of_its_own():
+    identity = _identity()
+    worker, probe = _worker(identity)
+    registration = NativeShardRegistration(
+        endpoint="http://native.example", model_id=identity.artifact.artifact_id, identity=identity,
+        worker=worker, probe=probe, device="cpu", capacity=ExecutionCapacity(1, 1, 1),
+    )
+    published: list[dict] = []
+    withdrawn: list[str] = []
+    registrar = NativeCapabilityRegistrar(registration, register=published.append, withdraw=withdrawn.append)
+    registrar.publish()
+    registrar.unavailable("worker exited")
+    assert published[0]["capability_report"]["backend"]["backend_id"] == identity.recipe.axes["backend_id"]
+    assert withdrawn == ["worker exited"]
+
+
+def test_native_registration_refuses_a_probe_for_a_different_range():
+    identity = _identity()
+    worker, probe = _worker(identity)
+    wrong = NativeWorkerProbe(**{**probe.__dict__, "shard_end": probe.shard_end + 1})
+    try:
+        NativeShardRegistration(
+            endpoint="http://native.example", model_id=identity.artifact.artifact_id, identity=identity,
+            worker=worker, probe=wrong, device="cpu", capacity=ExecutionCapacity(1, 1, 1),
+        )
+    except NativeRegistrationError:
+        return
+    raise AssertionError("different worker range must not register")