diff --git a/e2e/test_entities.py b/e2e/test_entities.py index a79a09b0d6..6bef8a3fac 100644 --- a/e2e/test_entities.py +++ b/e2e/test_entities.py @@ -20,7 +20,11 @@ import uuid import pytest -from nemo_platform import APIStatusError, NeMoPlatform +from nemo_platform import NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.client.errors import NemoHTTPError as APIStatusError +from nemo_platform_plugin.entities.client import EntitiesClient +from nemo_platform_plugin.entities.types import EntityCreateInput, EntityUpdate, ListEntitiesQueryParams from nmp.testing import as_service_for ENTITY_TYPE = "e2e-test-entity" @@ -66,57 +70,60 @@ def test_entity_crud_lifecycle(entity_store_sdk: NeMoPlatform, workspace: str): 4. Delete the entity 5. Verify it no longer exists """ + entities = client_from_platform(entity_store_sdk, EntitiesClient) entity_name = _unique_name() initial_data = {"key": "initial-value", "nested": {"field": 123}} # Create entity - entity = entity_store_sdk.entities.create( + entity = entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - name=entity_name, - data=initial_data, - ) + body=EntityCreateInput( + name=entity_name, + data=initial_data, + ), + ).data() assert entity.name == entity_name assert entity.workspace == workspace assert entity.entity_type == ENTITY_TYPE assert entity.data["key"] == "initial-value" - assert entity.data["nested"]["field"] == 123 # ty: ignore[not-subscriptable] + assert entity.data["nested"]["field"] == 123 try: # Retrieve by name - retrieved = entity_store_sdk.entities.get_entity_by_name( + retrieved = entities.get_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, - ) + ).data() assert retrieved.name == entity_name assert retrieved.id == entity.id assert retrieved.data == initial_data # Update entity updated_data = {"key": "updated-value", "nested": {"field": 456}, "new_field": True} - updated = entity_store_sdk.entities.update_entity_by_name( + updated = entities.update_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, - data=updated_data, - ) + body=EntityUpdate(data=updated_data), + ).data() assert updated.name == entity_name assert updated.data["key"] == "updated-value" - assert updated.data["nested"]["field"] == 456 # ty: ignore[not-subscriptable] + assert updated.data["nested"]["field"] == 456 assert updated.data["new_field"] is True # Verify update persisted - retrieved_after_update = entity_store_sdk.entities.get_entity_by_name( + retrieved_after_update = entities.get_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, - ) + ).data() assert retrieved_after_update.data == updated_data finally: # Delete entity - entity_store_sdk.entities.delete_entity_by_name( + entities.delete_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -124,7 +131,7 @@ def test_entity_crud_lifecycle(entity_store_sdk: NeMoPlatform, workspace: str): # Verify entity no longer exists with pytest.raises(APIStatusError) as exc_info: - entity_store_sdk.entities.get_entity_by_name( + entities.get_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -139,6 +146,7 @@ def test_entity_with_project(sdk: NeMoPlatform, entity_store_sdk: NeMoPlatform, CRUD uses service credentials. Verifies that entities can be associated with projects and that the association is persisted and retrievable. """ + entities = client_from_platform(entity_store_sdk, EntitiesClient) project_name = _unique_name("project") entity_name = _unique_name() @@ -152,26 +160,28 @@ def test_entity_with_project(sdk: NeMoPlatform, entity_store_sdk: NeMoPlatform, try: # Create entity within project - entity = entity_store_sdk.entities.create( + entity = entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - name=entity_name, - data={"project_data": "value"}, - project=project_name, - ) + body=EntityCreateInput( + name=entity_name, + data={"project_data": "value"}, + project=project_name, + ), + ).data() assert entity.name == entity_name assert entity.project == project_name # Retrieve and verify project association - retrieved = entity_store_sdk.entities.get_entity_by_name( + retrieved = entities.get_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, - ) + ).data() assert retrieved.project == project_name # Delete entity - entity_store_sdk.entities.delete_entity_by_name( + entities.delete_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -188,28 +198,31 @@ def test_entity_without_project(entity_store_sdk: NeMoPlatform, workspace: str): Verifies that entities can exist at the workspace level without being associated with any project. """ + entities = client_from_platform(entity_store_sdk, EntitiesClient) entity_name = _unique_name() - entity = entity_store_sdk.entities.create( + entity = entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - name=entity_name, - data={"standalone": True}, - ) + body=EntityCreateInput( + name=entity_name, + data={"standalone": True}, + ), + ).data() try: assert entity.name == entity_name assert entity.project is None - retrieved = entity_store_sdk.entities.get_entity_by_name( + retrieved = entities.get_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, - ) + ).data() assert retrieved.project is None finally: - entity_store_sdk.entities.delete_entity_by_name( + entities.delete_entity_by_name( name=entity_name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -224,62 +237,65 @@ def test_entity_list_and_sorting(entity_store_sdk: NeMoPlatform, workspace: str) 2. Sorting by created_at works (ascending and descending) 3. Sorting by name works """ + entities = client_from_platform(entity_store_sdk, EntitiesClient) entity_names = [_unique_name(f"sort-{i:02d}") for i in range(5)] created_entities = [] try: # Create entities in order for name in entity_names: - entity = entity_store_sdk.entities.create( + entity = entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - name=name, - data={"order": name}, - ) + body=EntityCreateInput( + name=name, + data={"order": name}, + ), + ).data() time.sleep(1) created_entities.append(entity) # List all entities of this type - response = entity_store_sdk.entities.list( + response = entities.list_entities( entity_type=ENTITY_TYPE, workspace=workspace, ) - listed_names = {e.name for e in response.data} + listed_names = {e.name for e in response.items()} for name in entity_names: assert name in listed_names # Test descending sort by created_at (default, newest first) - response_desc = entity_store_sdk.entities.list( + response_desc = entities.list_entities( entity_type=ENTITY_TYPE, workspace=workspace, - sort="-created_at", + query_params=ListEntitiesQueryParams(sort="-created_at"), ) - desc_names = [e.name for e in response_desc.data if e.name in entity_names] + desc_names = [e.name for e in response_desc.items() if e.name in entity_names] assert desc_names == list(reversed(entity_names)) # Test ascending sort by created_at (oldest first) - response_asc = entity_store_sdk.entities.list( + response_asc = entities.list_entities( entity_type=ENTITY_TYPE, workspace=workspace, - sort="created_at", + query_params=ListEntitiesQueryParams(sort="created_at"), ) - asc_names = [e.name for e in response_asc.data if e.name in entity_names] + asc_names = [e.name for e in response_asc.items() if e.name in entity_names] assert asc_names == entity_names # Test sort by name - response_by_name = entity_store_sdk.entities.list( + response_by_name = entities.list_entities( entity_type=ENTITY_TYPE, workspace=workspace, - sort="name", + query_params=ListEntitiesQueryParams(sort="name"), ) - name_sorted = [e.name for e in response_by_name.data if e.name in entity_names] + name_sorted = [e.name for e in response_by_name.items() if e.name in entity_names] assert name_sorted == sorted(entity_names) finally: # Clean up all created entities for name in entity_names: try: - entity_store_sdk.entities.delete_entity_by_name( + entities.delete_entity_by_name( name=name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -298,56 +314,63 @@ def test_entity_search_filter(entity_store_sdk: NeMoPlatform, workspace: str): entity_alpha = f"{prefix}-alpha" entity_beta = f"{prefix}-beta" + entities = client_from_platform(entity_store_sdk, EntitiesClient) try: # Create two entities with different data - entity_store_sdk.entities.create( + entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - name=entity_alpha, - data={"category": "alpha", "value": 100}, - ) - entity_store_sdk.entities.create( + body=EntityCreateInput( + name=entity_alpha, + data={"category": "alpha", "value": 100}, + ), + ).data() + entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - name=entity_beta, - data={"category": "beta", "value": 200}, - ) + body=EntityCreateInput( + name=entity_beta, + data={"category": "beta", "value": 200}, + ), + ).data() # Filter by exact name match filter_query = json.dumps({"name": {"$eq": entity_alpha}}) - response = entity_store_sdk.entities.list( + response = entities.list_entities( entity_type=ENTITY_TYPE, workspace=workspace, - filter=filter_query, + query_params=ListEntitiesQueryParams(filter=filter_query), ) - assert len(response.data) == 1 - assert response.data[0].name == entity_alpha + response_list = list(response.items()) + assert len(response_list) == 1 + assert response_list[0].name == entity_alpha # Filter by name pattern (like) filter_query = json.dumps({"name": {"$like": f"{prefix}%"}}) - response = entity_store_sdk.entities.list( + response = entities.list_entities( entity_type=ENTITY_TYPE, workspace=workspace, - filter=filter_query, + query_params=ListEntitiesQueryParams(filter=filter_query), ) - found_names = {e.name for e in response.data} + found_names = {e.name for e in response.items()} assert entity_alpha in found_names assert entity_beta in found_names # Filter by data field filter_query = json.dumps({"data.category": {"$eq": "beta"}}) - response = entity_store_sdk.entities.list( + response = entities.list_entities( entity_type=ENTITY_TYPE, workspace=workspace, - filter=filter_query, + query_params=ListEntitiesQueryParams(filter=filter_query), ) - assert len(response.data) == 1 - assert response.data[0].name == entity_beta + response_list = list(response.items()) + assert len(response_list) == 1 + assert response_list[0].name == entity_beta finally: for name in [entity_alpha, entity_beta]: try: - entity_store_sdk.entities.delete_entity_by_name( + entities.delete_entity_by_name( name=name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -362,31 +385,33 @@ def test_entity_rename(entity_store_sdk: NeMoPlatform, workspace: str): Verifies that entities can be renamed and the old name no longer works after rename. """ + entities = client_from_platform(entity_store_sdk, EntitiesClient) old_name = _unique_name("old") new_name = _unique_name("new") - entity = entity_store_sdk.entities.create( + entity = entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - name=old_name, - data={"test": "rename"}, - ) + body=EntityCreateInput( + name=old_name, + data={"test": "rename"}, + ), + ).data() try: # Rename entity - renamed = entity_store_sdk.entities.update_entity_by_name( + renamed = entities.update_entity_by_name( name=old_name, entity_type=ENTITY_TYPE, workspace=workspace, - data=entity.data, - new_name=new_name, - ) + body=EntityUpdate(data=entity.data, new_name=new_name), + ).data() assert renamed.name == new_name assert renamed.id == entity.id # Verify old name no longer works with pytest.raises(APIStatusError) as exc_info: - entity_store_sdk.entities.get_entity_by_name( + entities.get_entity_by_name( name=old_name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -394,17 +419,17 @@ def test_entity_rename(entity_store_sdk: NeMoPlatform, workspace: str): assert exc_info.value.status_code == 404 # Verify new name works - retrieved = entity_store_sdk.entities.get_entity_by_name( + retrieved = entities.get_entity_by_name( name=new_name, entity_type=ENTITY_TYPE, workspace=workspace, - ) + ).data() assert retrieved.name == new_name finally: # Clean up with new name try: - entity_store_sdk.entities.delete_entity_by_name( + entities.delete_entity_by_name( name=new_name, entity_type=ENTITY_TYPE, workspace=workspace, @@ -418,11 +443,14 @@ def test_entity_auto_generated_name(entity_store_sdk: NeMoPlatform, workspace: s When no name is provided, the API should auto-generate a unique name. """ - entity = entity_store_sdk.entities.create( + entities = client_from_platform(entity_store_sdk, EntitiesClient) + entity = entities.create_entity( entity_type=ENTITY_TYPE, workspace=workspace, - data={"auto_name": True}, - ) + body=EntityCreateInput( + data={"auto_name": True}, + ), + ).data() try: assert entity.name is not None @@ -431,15 +459,15 @@ def test_entity_auto_generated_name(entity_store_sdk: NeMoPlatform, workspace: s assert ENTITY_TYPE.replace("_", "-").replace("-", "") in entity.name.replace("-", "") or entity.name # Verify we can retrieve by the generated name - retrieved = entity_store_sdk.entities.get_entity_by_name( + retrieved = entities.get_entity_by_name( name=entity.name, entity_type=ENTITY_TYPE, workspace=workspace, - ) + ).data() assert retrieved.id == entity.id finally: - entity_store_sdk.entities.delete_entity_by_name( + entities.delete_entity_by_name( name=entity.name, entity_type=ENTITY_TYPE, workspace=workspace, diff --git a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/api/v2/events.py b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/api/v2/events.py index 31a8f5c2aa..5235627b87 100644 --- a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/api/v2/events.py +++ b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/api/v2/events.py @@ -24,6 +24,8 @@ from nemo_iron_swarm_plugin.entities import IRON_SWARM_RUN_TYPE from nemo_iron_swarm_plugin.filesets import download_fileset from nemo_platform_plugin.authz import CallerKind, path_rule +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.entities.client import EntitiesClient from pydantic import BaseModel from starlette.concurrency import run_in_threadpool @@ -153,10 +155,14 @@ def _fileset_fallback(workspace: str, name: str, stream: Any, after: int) -> lis sdk = _get_sdk() # get_entity_by_name returns a generic Entity — its domain fields live under `.data` # (same access pattern as sdk.py::_run_to_dict), not as top-level attributes. - run = sdk.entities.get_entity_by_name( - name=name, - entity_type=IRON_SWARM_RUN_TYPE, - workspace=workspace, + run = ( + client_from_platform(sdk, EntitiesClient) + .get_entity_by_name( + name=name, + entity_type=IRON_SWARM_RUN_TYPE, + workspace=workspace, + ) + .data() ) fileset_ref = (getattr(run, "data", None) or {}).get("events_fileset") if fileset_ref: diff --git a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/manifest.py b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/manifest.py index 724ec48a5e..e56cfdec89 100644 --- a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/manifest.py +++ b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/manifest.py @@ -19,6 +19,9 @@ from nemo_iron_swarm_plugin.entities import IRON_SWARM_MANIFEST_TYPE from nemo_iron_swarm_plugin.filesets import download_and_extract_project, upload_project_dir from nemo_iron_swarm_plugin.jobs.errors import CATEGORY_FILESET, CATEGORY_MANIFEST, IronSwarmRunError +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.entities.client import EntitiesClient +from nemo_platform_plugin.entities.types import EntityUpdate from nemo_platform_plugin.job_context import JobContext logger = logging.getLogger(__name__) @@ -147,8 +150,10 @@ def _materialize_manifest( raise IronSwarmRunError( CATEGORY_MANIFEST, "running a saved manifest requires the platform SDK (submit the job, don't run locally)." ) - record = sdk.entities.get_entity_by_name( - name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace + record = ( + client_from_platform(sdk, EntitiesClient) + .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace) + .data() ) data = {**(getattr(record, "data", {}) or {}), **(config_overrides or {})} manifest_dir = ctx.storage.persistent @@ -255,8 +260,11 @@ def _persist_upgraded_bundle(sdk: Any, manifest_id: str, ctx: JobContext, record updated = {**(getattr(record, "data", {}) or {})} updated["agent_fileset"] = fileset updated["manifest_yaml"] = yaml.safe_dump(resolved.manifest, sort_keys=False) - sdk.entities.update_entity_by_name( - name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace, data=updated + client_from_platform(sdk, EntitiesClient).update_entity_by_name( + name=manifest_id, + entity_type=IRON_SWARM_MANIFEST_TYPE, + workspace=ctx.workspace, + body=EntityUpdate(data=updated), ) except Exception: # the war-game matters more than the upgrade; it retries next run logger.warning("could not freeze manifest %s on this run", manifest_id, exc_info=True) diff --git a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/records.py b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/records.py index 13c1882ec6..432caa34af 100644 --- a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/records.py +++ b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/jobs/records.py @@ -21,6 +21,9 @@ ) from nemo_iron_swarm_plugin.jobs import benign_suite from nemo_iron_swarm_plugin.jobs.errors import RunFailure +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.entities.client import EntitiesClient +from nemo_platform_plugin.entities.types import EntityCreateInput, EntityUpdate from nemo_platform_plugin.entity_client import NemoEntitiesClient from nemo_platform_plugin.job_context import JobContext @@ -74,7 +77,11 @@ def _create_run(sdk: Any, *, workspace: str, data: dict[str, Any]) -> str | None if sdk is None or not hasattr(sdk, "entities"): return None try: - entity = sdk.entities.create(IRON_SWARM_RUN_TYPE, workspace=workspace, data=data) + entity = ( + client_from_platform(sdk, EntitiesClient) + .create_entity(entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace, body=EntityCreateInput(data=data)) + .data() + ) return getattr(entity, "name", None) except Exception: # recording is best-effort, not part of the war-game logger.warning("failed to persist IronSwarmRun record", exc_info=True) @@ -120,7 +127,11 @@ def _run_facts(sdk: Any, *, workspace: str, name: str) -> tuple[str, int]: if sdk is None or not hasattr(sdk, "entities"): return "", 0 try: - record = sdk.entities.get_entity_by_name(name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace) + record = ( + client_from_platform(sdk, EntitiesClient) + .get_entity_by_name(name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace) + .data() + ) data = getattr(record, "data", {}) or {} port = data.get("port") return str(data.get("agent") or ""), int(port) if isinstance(port, int) else 0 @@ -134,7 +145,9 @@ def _update_run(sdk: Any, *, workspace: str, name: str, data: dict[str, Any]) -> if sdk is None or not hasattr(sdk, "entities"): return try: - sdk.entities.update_entity_by_name(name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace, data=data) + client_from_platform(sdk, EntitiesClient).update_entity_by_name( + name=name, entity_type=IRON_SWARM_RUN_TYPE, workspace=workspace, body=EntityUpdate(data=data) + ).data() except Exception: # recording is best-effort, not part of the war-game logger.warning("failed to update IronSwarmRun record", exc_info=True) @@ -144,8 +157,10 @@ def _manifest_rounds(sdk: Any, manifest_id: str, ctx: JobContext) -> int: if sdk is None or not hasattr(sdk, "entities"): return 1 try: - record = sdk.entities.get_entity_by_name( - name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace + record = ( + client_from_platform(sdk, EntitiesClient) + .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace) + .data() ) rounds = (getattr(record, "data", {}) or {}).get("rounds") return rounds if isinstance(rounds, int) and rounds >= 1 else 1 @@ -159,8 +174,10 @@ def _manifest_models(sdk: Any, manifest_id: str, ctx: JobContext) -> dict[str, A if sdk is None or not hasattr(sdk, "entities"): return {} try: - record = sdk.entities.get_entity_by_name( - name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace + record = ( + client_from_platform(sdk, EntitiesClient) + .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace) + .data() ) models = (getattr(record, "data", {}) or {}).get("models") return models if isinstance(models, dict) else {} @@ -174,8 +191,10 @@ def _cached_benign_suite(sdk: Any, manifest_id: str, ctx: JobContext) -> list[di if sdk is None or not hasattr(sdk, "entities"): return [] try: - record = sdk.entities.get_entity_by_name( - name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace + record = ( + client_from_platform(sdk, EntitiesClient) + .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=ctx.workspace) + .data() ) suite = (getattr(record, "data", {}) or {}).get("benign_suite") or [] return [row for row in suite if isinstance(row, dict)] @@ -216,15 +235,17 @@ def _persist_benign_suite( if sdk is None or not hasattr(sdk, "entities") or not suite: return try: - record = sdk.entities.get_entity_by_name( - name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace + record = ( + client_from_platform(sdk, EntitiesClient) + .get_entity_by_name(name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace) + .data() ) data = dict(getattr(record, "data", {}) or {}) data["benign_suite"] = suite if interview: data["benign_interview"] = interview - sdk.entities.update_entity_by_name( - name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace, data=data + client_from_platform(sdk, EntitiesClient).update_entity_by_name( + name=manifest_id, entity_type=IRON_SWARM_MANIFEST_TYPE, workspace=workspace, body=EntityUpdate(data=data) ) except Exception: # caching is best-effort, not part of the war-game logger.warning("failed to cache benign suite on manifest %s", manifest_id, exc_info=True) diff --git a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/sdk.py b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/sdk.py index 301473714e..ca73c84720 100644 --- a/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/sdk.py +++ b/plugins/nemo-iron-swarm/src/nemo_iron_swarm_plugin/sdk.py @@ -25,6 +25,9 @@ from nemo_iron_swarm_plugin.jobs.run import IronSwarmRunJob from nemo_iron_swarm_plugin.jobs.synth_benign import IronSwarmSynthBenignJob from nemo_platform import AsyncNeMoPlatform, NeMoPlatform +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.entities.client import EntitiesClient +from nemo_platform_plugin.entities.types import ListEntitiesQueryParams from nemo_platform_plugin.scheduler import NemoJobScheduler from nemo_platform_plugin.sdk import NemoPluginSDKResources @@ -108,8 +111,16 @@ def _list_newest(platform: NeMoPlatform, entity_type: str, *, workspace: str, li ``page_size`` bounds the *page*, not the total — iterating it walks the entire history. We ask for one page of *limit* and take only that page's items, which is a single request. """ - page = platform.entities.list(entity_type, workspace=workspace, sort="-created_at", page_size=limit) - return [_run_to_dict(item) for item in itertools.islice(page, limit)] + page = ( + client_from_platform(platform, EntitiesClient) + .list_entities( + entity_type=entity_type, + workspace=workspace, + query_params=ListEntitiesQueryParams(sort="-created_at", page_size=limit), + ) + .page() + ) + return [_run_to_dict(item) for item in itertools.islice(page.items, limit)] class _RunsResource: diff --git a/plugins/nemo-iron-swarm/tests/unit/conftest.py b/plugins/nemo-iron-swarm/tests/unit/conftest.py new file mode 100644 index 0000000000..36410a5df2 --- /dev/null +++ b/plugins/nemo-iron-swarm/tests/unit/conftest.py @@ -0,0 +1,74 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Shared test doubles for the iron-swarm plugin unit tests. + +The records/manifest/sdk/events modules talk to the entity store through the typed +:class:`EntitiesClient` via :func:`client_from_platform`. Unit tests fake the entity store +as a ``SimpleNamespace(entities=...)`` shape, so ``client_from_platform`` must be patched +at each consuming module's boundary to route the typed-client calls back onto that fake. +This keeps the per-test ``entities`` fakes (and the capturing run-service doubles) unchanged. + +The typed client method calls are translated to the fake's shape here, once: +create_entity/get_entity_by_name/update_entity_by_name/list_entities -> entities.create/ +get_entity_by_name/update_entity_by_name/list, with the request ``body``/``query_params`` +flattened back to the ``data``/``sort``/``page_size`` kwargs the fakes expect. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from typing import Any + +import pytest +from nemo_iron_swarm_plugin import sdk as sdk_module +from nemo_iron_swarm_plugin.api.v2 import events as events_module +from nemo_iron_swarm_plugin.jobs import manifest as manifest_module +from nemo_iron_swarm_plugin.jobs import records as records_module + + +def _data(body: Any) -> Any: + """Flatten the request body (EntityCreateInput/EntityUpdate) to its data dict.""" + return body.data + + +def _ok(value: Any) -> Any: + """Typed-client responses expose ``.data()``; wrap the fake so call sites can call it.""" + return SimpleNamespace(data=lambda: value) + + +def _page(value: Any) -> Any: + """Typed paginated responses expose ``.page().items``; surface the fake's iteration through it.""" + return SimpleNamespace(page=lambda: SimpleNamespace(items=iter(value))) + + +def _fake_entities_client(platform: Any, _client_cls: Any) -> Any: + """Build a typed-client-shaped stub that delegates to the fake ``platform.entities`` namespace.""" + entities = platform.entities + return SimpleNamespace( + create_entity=lambda *, entity_type, workspace, body: _ok( + entities.create(entity_type, workspace=workspace, data=_data(body)) + ), + get_entity_by_name=lambda *, name, entity_type, workspace: _ok( + entities.get_entity_by_name(name=name, entity_type=entity_type, workspace=workspace) + ), + update_entity_by_name=lambda *, name, entity_type, workspace, body: _ok( + entities.update_entity_by_name(name=name, entity_type=entity_type, workspace=workspace, data=_data(body)) + ), + list_entities=lambda *, entity_type, workspace, query_params=None: _page( + entities.list( + entity_type, + sort=query_params.get("sort") if query_params else None, + workspace=workspace, + page_size=query_params.get("page_size") if query_params else None, + ) + ), + ) + + +@pytest.fixture(autouse=True) +def _fake_entities_client_boundary(monkeypatch: pytest.MonkeyPatch) -> None: + """Route typed-client entity calls onto the fake ``entities`` namespace in every consuming module.""" + for mod in (records_module, manifest_module, events_module): + monkeypatch.setattr(mod, "client_from_platform", _fake_entities_client) + monkeypatch.setattr(sdk_module, "client_from_platform", _fake_entities_client) diff --git a/services/core/entities/tests/integration/test_workspaces_crud_with_auth.py b/services/core/entities/tests/integration/test_workspaces_crud_with_auth.py index 8d62390f18..d94fbd2880 100644 --- a/services/core/entities/tests/integration/test_workspaces_crud_with_auth.py +++ b/services/core/entities/tests/integration/test_workspaces_crud_with_auth.py @@ -19,6 +19,9 @@ import pytest from nemo_platform import ConflictError, NeMoPlatform, PermissionDeniedError +from nemo_platform_plugin.client.adapter import client_from_platform +from nemo_platform_plugin.entities.client import EntitiesClient +from nemo_platform_plugin.entities.types import EntityCreateInput from nmp.core.entities.service import EntitiesService from nmp.testing import TEST_USER_EMAIL, create_test_client, short_unique_name @@ -560,12 +563,14 @@ def test_delete_workspace_with_entities_marks_for_deletion(self, sdk: NeMoPlatfo # Create entity as service principal (generic entities API requires service credentials) with as_service(sdk, "entities"): - sdk.entities.create( + client_from_platform(sdk, EntitiesClient).create_entity( workspace=workspace_name, entity_type="test-entity-type", - name="test-entity", - data={"key": "value"}, - ) + body=EntityCreateInput( + name="test-entity", + data={"key": "value"}, + ), + ).data() with as_user(sdk, admin_email): sdk.workspaces.delete(workspace_name) diff --git a/tests/agentic-use/entities-basic-cli-easy/tests/test_outputs.py b/tests/agentic-use/entities-basic-cli-easy/tests/test_outputs.py index 2bca7ee124..ccddb6a0bf 100644 --- a/tests/agentic-use/entities-basic-cli-easy/tests/test_outputs.py +++ b/tests/agentic-use/entities-basic-cli-easy/tests/test_outputs.py @@ -14,29 +14,29 @@ import os import pytest -from nemo_platform import NeMoPlatform +from nemo_platform_plugin.entities.client import EntitiesClient WORKSPACE = "default" @pytest.fixture -def client() -> NeMoPlatform: +def client() -> EntitiesClient: nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080") - return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) + return EntitiesClient(base_url=nmp_base_url, workspace=WORKSPACE) -def test_harbor_test_model_deleted(client: NeMoPlatform) -> None: +def test_harbor_test_model_deleted(client: EntitiesClient) -> None: """Test that harbor-test-model was deleted after CRUD operations.""" - response = client.entities.list(entity_type="model") - entity_names = [e.name for e in response.data] + response = client.list_entities(entity_type="model") + entity_names = [e.name for e in response.page().items] assert "harbor-test-model" not in entity_names, ( f"Entity 'harbor-test-model' should have been deleted but still exists! Found: {entity_names}" ) -def test_harbor_final_dataset_exists(client: NeMoPlatform) -> None: +def test_harbor_final_dataset_exists(client: EntitiesClient) -> None: """Test that harbor-final-dataset was created and has correct data.""" - response = client.entities.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset") + response = client.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset").data() assert response.name == "harbor-final-dataset", ( f"Expected entity name 'harbor-final-dataset', got '{response.name}'" ) diff --git a/tests/agentic-use/entities-basic-cli/tests/test_outputs.py b/tests/agentic-use/entities-basic-cli/tests/test_outputs.py index 2bca7ee124..ccddb6a0bf 100644 --- a/tests/agentic-use/entities-basic-cli/tests/test_outputs.py +++ b/tests/agentic-use/entities-basic-cli/tests/test_outputs.py @@ -14,29 +14,29 @@ import os import pytest -from nemo_platform import NeMoPlatform +from nemo_platform_plugin.entities.client import EntitiesClient WORKSPACE = "default" @pytest.fixture -def client() -> NeMoPlatform: +def client() -> EntitiesClient: nmp_base_url = os.environ.get("NMP_BASE_URL", "http://localhost:8080") - return NeMoPlatform(base_url=nmp_base_url, workspace=WORKSPACE) + return EntitiesClient(base_url=nmp_base_url, workspace=WORKSPACE) -def test_harbor_test_model_deleted(client: NeMoPlatform) -> None: +def test_harbor_test_model_deleted(client: EntitiesClient) -> None: """Test that harbor-test-model was deleted after CRUD operations.""" - response = client.entities.list(entity_type="model") - entity_names = [e.name for e in response.data] + response = client.list_entities(entity_type="model") + entity_names = [e.name for e in response.page().items] assert "harbor-test-model" not in entity_names, ( f"Entity 'harbor-test-model' should have been deleted but still exists! Found: {entity_names}" ) -def test_harbor_final_dataset_exists(client: NeMoPlatform) -> None: +def test_harbor_final_dataset_exists(client: EntitiesClient) -> None: """Test that harbor-final-dataset was created and has correct data.""" - response = client.entities.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset") + response = client.get_entity_by_name(name="harbor-final-dataset", entity_type="dataset").data() assert response.name == "harbor-final-dataset", ( f"Expected entity name 'harbor-final-dataset', got '{response.name}'" )