Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,9 @@
from nemo_platform_plugin.entities.base import (
EntityTypeLike as EntityTypeLike,
)
from nemo_platform_plugin.entities.base import (
EntityUpdateClientProtocol as EntityUpdateClientProtocol,
)
from nemo_platform_plugin.entities.base import (
EntityValidationError as EntityValidationError,
)
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -231,9 +231,21 @@ class EntityToken(Protocol):


class EntityGetterProtocol(Protocol[EntityT]):
"""Protocol for entity clients that can fetch entities by workspace/name."""
"""Protocol for entity clients that can fetch entities by workspace/name.

async def get(self, entity_type: Type[EntityT], *, name: str, workspace: str) -> EntityT: ...
``parent`` addresses a **child** entity, which is unique within
``(workspace, entity_type, parent, name)`` rather than by name alone. It is optional, so
fetching a root entity is unchanged.
"""

async def get(
self,
entity_type: Type[EntityT],
*,
name: str,
workspace: str,
parent: Optional[str] = None,
) -> EntityT: ...


class EntityDeleteClientProtocol(EntityGetterProtocol[EntityT], Protocol[EntityT]):
Expand Down Expand Up @@ -266,6 +278,20 @@ class EntityClientProtocol(EntityDeleteClientProtocol[EntityT], Protocol[EntityT
async def create(self, entity: EntityT) -> EntityT: ...


class EntityUpdateClientProtocol(Protocol[EntityT]):
"""Protocol for entity clients that can update an existing entity.

Separate from :class:`EntityClientProtocol` rather than folded into it: ``update`` is a
read-modify-write against the ``db_version`` optimistic lock, and most services never need it.
Compose it with the CRUD protocol where a service does::

class Store(EntityClientProtocol[MyEntity], EntityUpdateClientProtocol[MyEntity], Protocol):
...
"""

async def update(self, entity: EntityT, *, original_name: Optional[str] = None) -> EntityT: ...


class AnyEntityGetterProtocol(Protocol):
"""Protocol for clients that can fetch any entity model type."""

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,9 @@
from nemo_platform_plugin.entities import (
EntityNotFoundError as NemoEntityNotFoundError,
)
from nemo_platform_plugin.entities import (
EntityUpdateClientProtocol as NemoEntityUpdateClientProtocol,
)
from nemo_platform_plugin.entities import (
EntityValidationError as NemoEntityValidationError,
)
Expand All @@ -65,6 +68,7 @@
"NemoEntitiesClient",
"NemoEntitiesClientProtocol",
"NemoAnyEntityDeleteClientProtocol",
"NemoEntityUpdateClientProtocol",
"NemoAnyEntityGetterProtocol",
"NemoEntityDeleteClientProtocol",
"NemoEntityGetterProtocol",
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,80 @@
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

"""The entity-client protocols must describe the client they stand in for.

A protocol that has drifted from its implementation is worse than no protocol: a service typed
against it either fails type-checking on correct code, or type-checks against a method the real
client does not have. These tests pin that relationship so the two cannot silently diverge.
"""

from __future__ import annotations

import inspect
from typing import Protocol, TypeVar

from nemo_platform_plugin.entities import (
EntityBase,
EntityClient,
EntityClientProtocol,
EntityGetterProtocol,
EntityUpdateClientProtocol,
)

EntityT = TypeVar("EntityT", bound=EntityBase)


class _Entity(EntityBase):
__entity_type__ = "protocol_conformance_probe"


class _ReadWriteStore(
EntityClientProtocol[EntityT],
EntityUpdateClientProtocol[EntityT],
Protocol[EntityT],
):
"""The shape a service needing the wider surface composes for itself.

Exists here to prove the pieces *compose*: ``update`` is a separate protocol precisely so a
service can opt into it alongside the CRUD one, rather than declaring a private protocol that
restates the whole surface.
"""


def _static_conformance(client: EntityClient) -> _ReadWriteStore[_Entity]:
"""Static assertion, checked by ``ty`` rather than at runtime.

If ``EntityClient`` ever stops satisfying the composed protocols — a renamed method, a changed
signature — this return fails type-checking. The runtime tests below document *which* parts
matter and why; this is what actually catches drift, because structural conformance is a
type-level property no ``hasattr`` check can verify.
"""
return client


def _signature(owner: object, method: str) -> inspect.Signature:
return inspect.signature(getattr(owner, method))


def test_update_is_its_own_protocol() -> None:
"""``update`` is opt-in. Most services never modify an entity in place, and folding ``update``
into the CRUD protocol would force each of them — and every one of their test doubles — to
satisfy a method they do not use."""
assert hasattr(EntityUpdateClientProtocol, "update")
assert not hasattr(EntityClientProtocol, "update")


def test_update_protocol_matches_the_client() -> None:
protocol = _signature(EntityUpdateClientProtocol, "update").parameters
client = _signature(EntityClient, "update").parameters
assert set(protocol) == set(client)
assert protocol["original_name"].kind is inspect.Parameter.KEYWORD_ONLY


def test_getter_accepts_parent_for_child_entities() -> None:
"""Child records are unique within ``(workspace, entity_type, parent, name)``, so the parent is
part of their address. It is optional, so fetching a root entity is unchanged."""
getter = _signature(EntityGetterProtocol, "get").parameters
assert "parent" in getter
assert getter["parent"].default is None
assert "parent" in _signature(EntityClient, "get").parameters
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,9 @@ def __init__(self) -> None:
self.delete_error: Exception | None = None
self.list_filter_operations: list[FilterOperation | None] = []

async def get(self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str) -> MetricBundleEntity:
async def get(
self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str, parent: str | None = None
) -> MetricBundleEntity:
key = (workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,9 @@ def seed(self, entity: _ResultEntityT) -> _ResultEntityT:
self.entities[(entity.__entity_type__, entity.workspace, entity.name)] = entity
return entity

async def get(self, entity_type: type[_ResultEntityT], *, workspace: str, name: str) -> _ResultEntityT:
async def get(
self, entity_type: type[_ResultEntityT], *, workspace: str, name: str, parent: str | None = None
) -> _ResultEntityT:
key = (entity_type.__entity_type__, workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,9 @@ async def create(self, entity: TaskEntity) -> TaskEntity:
self.entities[key] = entity
return entity

async def get(self, entity_type: type[TaskEntity], *, workspace: str, name: str) -> TaskEntity:
async def get(
self, entity_type: type[TaskEntity], *, workspace: str, name: str, parent: str | None = None
) -> TaskEntity:
key = (entity_type.__entity_type__, workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,9 @@ async def create(self, entity: TasksetEntity) -> TasksetEntity:
self.entities[key] = entity
return entity

async def get(self, entity_type: type[TasksetEntity], *, workspace: str, name: str) -> TasksetEntity:
async def get(
self, entity_type: type[TasksetEntity], *, workspace: str, name: str, parent: str | None = None
) -> TasksetEntity:
key = (entity_type.__entity_type__, workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -78,7 +78,9 @@ def __init__(self) -> None:
self.bump_version_on_next_delete = False
self.delete_expected_db_versions: list[int | None] = []

async def get(self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str) -> MetricBundleEntity:
async def get(
self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str, parent: str | None = None
) -> MetricBundleEntity:
key = (workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@ def seed(self, entity: _ResultEntityT) -> _ResultEntityT:
self.entities[(entity.__entity_type__, entity.workspace, entity.name)] = entity
return entity

async def get(self, entity_type: type[_ResultEntityT], *, workspace: str, name: str) -> _ResultEntityT:
async def get(
self, entity_type: type[_ResultEntityT], *, workspace: str, name: str, parent: str | None = None
) -> _ResultEntityT:
key = (entity_type.__entity_type__, workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
4 changes: 3 additions & 1 deletion plugins/nemo-evaluator/tests/api/v2/test_tasks_routes.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@ async def create(self, entity: TaskEntity) -> TaskEntity:
self.entities[key] = entity
return entity

async def get(self, entity_type: type[TaskEntity], *, workspace: str, name: str) -> TaskEntity:
async def get(
self, entity_type: type[TaskEntity], *, workspace: str, name: str, parent: str | None = None
) -> TaskEntity:
key = (entity_type.__entity_type__, workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@ async def create(self, entity: TasksetEntity) -> TasksetEntity:
self.entities[key] = entity
return entity

async def get(self, entity_type: type[TasksetEntity], *, workspace: str, name: str) -> TasksetEntity:
async def get(
self, entity_type: type[TasksetEntity], *, workspace: str, name: str, parent: str | None = None
) -> TasksetEntity:
key = (entity_type.__entity_type__, workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
4 changes: 3 additions & 1 deletion plugins/nemo-evaluator/tests/test_metric_refs.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,7 +60,9 @@ class _FakeEntityClient:
def __init__(self) -> None:
self.entities: dict[tuple[str, str], MetricBundleEntity] = {}

async def get(self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str) -> MetricBundleEntity:
async def get(
self, entity_type: type[MetricBundleEntity], *, workspace: str, name: str, parent: str | None = None
) -> MetricBundleEntity:
try:
return self.entities[(workspace, name)]
except KeyError:
Expand Down
4 changes: 3 additions & 1 deletion plugins/nemo-evaluator/tests/test_task_refs.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,9 @@ def __init__(self) -> None:
def add(self, entity: EntityBase) -> None:
self.entities[(entity.__entity_type__, entity.workspace, entity.name)] = entity

async def get(self, entity_type: type[_EntityT], *, workspace: str, name: str) -> _EntityT:
async def get(
self, entity_type: type[_EntityT], *, workspace: str, name: str, parent: str | None = None
) -> _EntityT:
key = (entity_type.__entity_type__, workspace, name)
if key not in self.entities:
raise NemoEntityNotFoundError(f"{workspace}/{name} not found")
Expand Down
Loading