From 14a9e57e77f74967c92074a674cf2807acf02226 Mon Sep 17 00:00:00 2001 From: spuentesp Date: Mon, 5 Jan 2026 10:05:37 -0300 Subject: [PATCH 1/2] feat(data-layer): DL-14 - Manage Relationships & State Tags MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements comprehensive relationship and state tag management in Neo4j for entity-to-entity connections and dynamic entity status tracking. ## Relationships (Neo4j edges) New schemas (relationships.py): - RelationshipType enum: 7 types (MEMBER_OF, OWNS, KNOWS, ALLIED_WITH, HOSTILE_TO, LOCATED_IN, PARTICIPATES_IN) - Direction enum: OUTGOING, INCOMING, BOTH for query filtering - RelationshipCreate/Update/Response: Full CRUD with property storage - RelationshipFilter/ListResponse: Flexible querying by entity, type, direction Neo4j tools (neo4j_tools.py): - neo4j_create_relationship: Create typed edges with entity validation - neo4j_get_relationship: Retrieve by Neo4j internal ID - neo4j_list_relationships: Filter by entity/type/direction with pagination - neo4j_update_relationship: Update properties, preserve created_at - neo4j_delete_relationship: Remove relationship edge ## State Tags (Dynamic entity status) New schemas (relationships.py): - StateTag enum: 16 tags for entity instance status - Vital: alive, dead, unconscious, wounded - Visibility: hidden, revealed - Disposition: hostile, friendly, neutral - Combat: prone, grappled, restrained, incapacitated - Mental: charmed, frightened, stunned, confused - StateTagUpdate/Response: Atomic add/remove operations Neo4j tools (neo4j_tools.py): - neo4j_update_state_tags: Atomic tag updates with archetype validation - neo4j_get_state_tags: Retrieve current tags for entity ## Base Schemas (base.py) Exported existing enums from relationships for external use ## Authority Rules (auth.py) Added 7 authority rules: - Relationship writes: CanonKeeper only - Relationship reads: All agents - State tag writes: CanonKeeper only - State tag reads: All agents ## Tests Added 23 comprehensive tests (test_relationship_tools.py): - Relationship CRUD: create (entity validation, all types), get, list (direction/type filtering), update, delete - State tags: add, remove, atomic add+remove, archetype validation, get - Error handling: entity not found, relationship not found, archetype validation - All tests use proper mocking patterns with @patch decorators ## Architecture Decisions - Relationships as Neo4j edges enable graph traversal and rich queries - State tags only on instances (not archetypes) enforces instance-specific state - Atomic tag updates prevent race conditions in concurrent modifications - Direction enum supports flexible relationship queries (incoming/outgoing/both) - All write operations require CanonKeeper authority for data integrity Implements: DL-14 Dependencies: DL-2 (Entity Management) Enables: DL-15 (Parties), DL-25 (Combat), gameplay mechanics All 23 tests passing ✅ 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 --- .../src/monitor_data/middleware/auth.py | 8 +- .../src/monitor_data/schemas/__init__.py | 23 + .../src/monitor_data/schemas/relationships.py | 172 +++++ .../src/monitor_data/tools/neo4j_tools.py | 406 +++++++++++ .../test_tools/test_relationship_tools.py | 690 ++++++++++++++++++ 5 files changed, 1297 insertions(+), 2 deletions(-) create mode 100644 packages/data-layer/src/monitor_data/schemas/relationships.py create mode 100644 packages/data-layer/tests/test_tools/test_relationship_tools.py diff --git a/packages/data-layer/src/monitor_data/middleware/auth.py b/packages/data-layer/src/monitor_data/middleware/auth.py index a11b6bda..72448511 100644 --- a/packages/data-layer/src/monitor_data/middleware/auth.py +++ b/packages/data-layer/src/monitor_data/middleware/auth.py @@ -44,11 +44,15 @@ "neo4j_delete_entity": ["CanonKeeper"], "neo4j_set_state_tags": ["CanonKeeper"], # ========================================================================= - # NEO4J OPERATIONS - Relationships + # NEO4J OPERATIONS - Relationships (DL-14) # ========================================================================= "neo4j_create_relationship": ["CanonKeeper"], - "neo4j_get_relationships": ["*"], + "neo4j_get_relationship": ["*"], + "neo4j_list_relationships": ["*"], + "neo4j_update_relationship": ["CanonKeeper"], "neo4j_delete_relationship": ["CanonKeeper"], + "neo4j_update_state_tags": ["CanonKeeper"], + "neo4j_get_state_tags": ["*"], # ========================================================================= # NEO4J OPERATIONS - Facts & Events # ========================================================================= diff --git a/packages/data-layer/src/monitor_data/schemas/__init__.py b/packages/data-layer/src/monitor_data/schemas/__init__.py index 07dad1ec..00f5fe1a 100644 --- a/packages/data-layer/src/monitor_data/schemas/__init__.py +++ b/packages/data-layer/src/monitor_data/schemas/__init__.py @@ -66,6 +66,18 @@ ResolutionFilter, ResolutionListResponse, ) +from monitor_data.schemas.relationships import ( + RelationshipType, + Direction, + StateTag, + RelationshipCreate, + RelationshipUpdate, + RelationshipResponse, + RelationshipFilter, + RelationshipListResponse, + StateTagUpdate, + StateTagResponse, +) # from monitor_data.schemas.entities import * # from monitor_data.schemas.facts import * @@ -118,4 +130,15 @@ "ResolutionResponse", "ResolutionFilter", "ResolutionListResponse", + # Relationship schemas + "RelationshipType", + "Direction", + "StateTag", + "RelationshipCreate", + "RelationshipUpdate", + "RelationshipResponse", + "RelationshipFilter", + "RelationshipListResponse", + "StateTagUpdate", + "StateTagResponse", ] diff --git a/packages/data-layer/src/monitor_data/schemas/relationships.py b/packages/data-layer/src/monitor_data/schemas/relationships.py new file mode 100644 index 00000000..8e01dade --- /dev/null +++ b/packages/data-layer/src/monitor_data/schemas/relationships.py @@ -0,0 +1,172 @@ +""" +Pydantic schemas for Relationship and State Tag operations (DL-14). + +LAYER: 1 (data-layer) +IMPORTS FROM: External libraries (pydantic, uuid, datetime, enum) and base schemas +CALLED BY: neo4j_tools.py + +These schemas define the data contracts for managing relationships between entities +and dynamic state tags on entity instances. Relationships are Neo4j edges with +typed connections and metadata. State tags track dynamic entity status. +""" + +from datetime import datetime +from enum import Enum +from typing import Optional, List, Dict, Any +from uuid import UUID + +from pydantic import BaseModel, Field + + +# ============================================================================= +# ENUMS +# ============================================================================= + + +class RelationshipType(str, Enum): + """Type of relationship between entities.""" + + MEMBER_OF = "MEMBER_OF" # Entity belongs to organization/group + OWNS = "OWNS" # Entity owns another entity/object + KNOWS = "KNOWS" # Social relationship - acquaintance + ALLIED_WITH = "ALLIED_WITH" # Formal alliance relationship + HOSTILE_TO = "HOSTILE_TO" # Antagonistic relationship + LOCATED_IN = "LOCATED_IN" # Spatial containment + PARTICIPATES_IN = "PARTICIPATES_IN" # Event/activity participation + + +class Direction(str, Enum): + """Direction for relationship queries.""" + + OUTGOING = "outgoing" # Relationships from entity to others + INCOMING = "incoming" # Relationships from others to entity + BOTH = "both" # Relationships in both directions + + +class StateTag(str, Enum): + """Dynamic state tags for entity instances.""" + + # Vital status + ALIVE = "alive" + DEAD = "dead" + UNCONSCIOUS = "unconscious" + WOUNDED = "wounded" + + # Visibility + HIDDEN = "hidden" + REVEALED = "revealed" + + # Disposition + HOSTILE = "hostile" + FRIENDLY = "friendly" + NEUTRAL = "neutral" + + # Combat states + PRONE = "prone" + GRAPPLED = "grappled" + RESTRAINED = "restrained" + INCAPACITATED = "incapacitated" + + # Mental states + CHARMED = "charmed" + FRIGHTENED = "frightened" + STUNNED = "stunned" + CONFUSED = "confused" + + +# ============================================================================= +# RELATIONSHIP CRUD SCHEMAS +# ============================================================================= + + +class RelationshipCreate(BaseModel): + """Request to create a relationship between entities.""" + + from_entity_id: UUID = Field(description="Source entity ID") + to_entity_id: UUID = Field(description="Target entity ID") + rel_type: RelationshipType + properties: Dict[str, Any] = Field( + default_factory=dict, + description="Optional properties (since, strength, notes, etc.)", + ) + + +class RelationshipUpdate(BaseModel): + """Request to update a relationship's properties.""" + + properties: Dict[str, Any] = Field( + description="Updated properties (replaces existing)" + ) + + +class RelationshipResponse(BaseModel): + """Response with relationship data.""" + + relationship_id: str = Field(description="Neo4j internal relationship ID") + from_entity_id: UUID + to_entity_id: UUID + rel_type: RelationshipType + properties: Dict[str, Any] + created_at: Optional[datetime] = Field( + None, description="When relationship was created" + ) + + model_config = {"from_attributes": True} + + +# ============================================================================= +# RELATIONSHIP QUERY SCHEMAS +# ============================================================================= + + +class RelationshipFilter(BaseModel): + """Filter parameters for listing relationships.""" + + entity_id: Optional[UUID] = Field( + None, description="Filter by entity (as source or target)" + ) + rel_type: Optional[RelationshipType] = Field( + None, description="Filter by relationship type" + ) + direction: Direction = Field( + default=Direction.BOTH, description="Direction relative to entity_id" + ) + limit: int = Field(default=50, ge=1, le=100) + offset: int = Field(default=0, ge=0) + + +class RelationshipListResponse(BaseModel): + """Response for list operations.""" + + relationships: List[RelationshipResponse] + total: int + limit: int + offset: int + + +# ============================================================================= +# STATE TAG SCHEMAS +# ============================================================================= + + +class StateTagUpdate(BaseModel): + """Request to update state tags on an entity instance.""" + + entity_id: UUID + add_tags: List[StateTag] = Field( + default_factory=list, description="Tags to add to entity" + ) + remove_tags: List[StateTag] = Field( + default_factory=list, description="Tags to remove from entity" + ) + + +class StateTagResponse(BaseModel): + """Response with entity's current state tags.""" + + entity_id: UUID + state_tags: List[StateTag] = Field( + default_factory=list, description="Current state tags on entity" + ) + + model_config = {"from_attributes": True} diff --git a/packages/data-layer/src/monitor_data/tools/neo4j_tools.py b/packages/data-layer/src/monitor_data/tools/neo4j_tools.py index 57103253..80736c85 100644 --- a/packages/data-layer/src/monitor_data/tools/neo4j_tools.py +++ b/packages/data-layer/src/monitor_data/tools/neo4j_tools.py @@ -65,6 +65,16 @@ RemovePartyMember, SetActivePC, ) +from monitor_data.schemas.relationships import ( + RelationshipCreate, + RelationshipUpdate, + RelationshipResponse, + RelationshipFilter, + RelationshipListResponse, + StateTagUpdate, + StateTagResponse, + Direction, +) # ============================================================================= @@ -3542,3 +3552,399 @@ def neo4j_delete_party(party_id: UUID) -> Dict[str, Any]: "party_id": str(party_id), "deleted_count": result[0]["deleted_count"] if result else 0, } + + +# ============================================================================= +# RELATIONSHIP TOOLS (DL-14) +# ============================================================================= + + +def neo4j_create_relationship(params: RelationshipCreate) -> RelationshipResponse: + """ + Create a typed relationship (edge) between two entities. + + Authority: CanonKeeper only + Use Case: DL-14 + + Args: + params: Relationship creation parameters + + Returns: + RelationshipResponse with created relationship data + + Raises: + ValueError: If either entity doesn't exist + """ + client = get_neo4j_client() + + # Validate both entities exist + from_exists = client.execute_read( + "MATCH (e:Entity {id: $entity_id}) RETURN e.id", + {"entity_id": str(params.from_entity_id)}, + ) + if not from_exists: + raise ValueError(f"From entity {params.from_entity_id} not found") + + to_exists = client.execute_read( + "MATCH (e:Entity {id: $entity_id}) RETURN e.id", + {"entity_id": str(params.to_entity_id)}, + ) + if not to_exists: + raise ValueError(f"To entity {params.to_entity_id} not found") + + # Create relationship with properties + now = datetime.now(timezone.utc) + props = {**params.properties, "created_at": now.isoformat()} + + create_query = f""" + MATCH (from:Entity {{id: $from_id}}) + MATCH (to:Entity {{id: $to_id}}) + CREATE (from)-[r:{params.rel_type.value} $props]->(to) + RETURN id(r) as rel_id, type(r) as rel_type, properties(r) as props + """ + + result = client.execute_write( + create_query, + { + "from_id": str(params.from_entity_id), + "to_id": str(params.to_entity_id), + "props": props, + }, + ) + + if not result: + raise ValueError("Failed to create relationship") + + rel_data = result[0] + return RelationshipResponse( + relationship_id=str(rel_data["rel_id"]), + from_entity_id=params.from_entity_id, + to_entity_id=params.to_entity_id, + rel_type=params.rel_type, + properties=rel_data["props"], + created_at=now, + ) + + +def neo4j_get_relationship(relationship_id: str) -> Optional[RelationshipResponse]: + """ + Get a relationship by its Neo4j internal ID. + + Authority: All agents + Use Case: DL-14 + + Args: + relationship_id: Neo4j relationship ID + + Returns: + RelationshipResponse if found, None otherwise + """ + client = get_neo4j_client() + + query = """ + MATCH (from:Entity)-[r]->(to:Entity) + WHERE id(r) = $rel_id + RETURN id(r) as rel_id, from.id as from_id, to.id as to_id, + type(r) as rel_type, properties(r) as props + """ + + result = client.execute_read(query, {"rel_id": int(relationship_id)}) + + if not result: + return None + + rel = result[0] + return RelationshipResponse( + relationship_id=str(rel["rel_id"]), + from_entity_id=UUID(rel["from_id"]), + to_entity_id=UUID(rel["to_id"]), + rel_type=rel["rel_type"], + properties=rel["props"], + created_at=( + datetime.fromisoformat(rel["props"].get("created_at")) + if rel["props"].get("created_at") + else None + ), + ) + + +def neo4j_list_relationships( + params: RelationshipFilter, +) -> RelationshipListResponse: + """ + List relationships with optional filtering. + + Authority: All agents + Use Case: DL-14 + + Args: + params: Filter parameters + + Returns: + RelationshipListResponse with matching relationships + """ + client = get_neo4j_client() + + # Build query based on filters + match_clause = "MATCH (from:Entity)-[r]->(to:Entity)" + where_clauses = [] + query_params: Dict[str, Any] = { + "limit": params.limit, + "offset": params.offset, + } + + if params.entity_id: + if params.direction == Direction.OUTGOING: + where_clauses.append("from.id = $entity_id") + elif params.direction == Direction.INCOMING: + where_clauses.append("to.id = $entity_id") + else: # BOTH + where_clauses.append("(from.id = $entity_id OR to.id = $entity_id)") + query_params["entity_id"] = str(params.entity_id) + + if params.rel_type: + where_clauses.append(f"type(r) = '{params.rel_type.value}'") + + where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else "" + + # Count query + count_query = f""" + {match_clause} + {where_clause} + RETURN count(r) as total + """ + + count_result = client.execute_read(count_query, query_params) + total = count_result[0]["total"] if count_result else 0 + + # Data query + data_query = f""" + {match_clause} + {where_clause} + RETURN id(r) as rel_id, from.id as from_id, to.id as to_id, + type(r) as rel_type, properties(r) as props + ORDER BY id(r) + SKIP $offset + LIMIT $limit + """ + + results = client.execute_read(data_query, query_params) + + relationships = [] + for rel in results: + relationships.append( + RelationshipResponse( + relationship_id=str(rel["rel_id"]), + from_entity_id=UUID(rel["from_id"]), + to_entity_id=UUID(rel["to_id"]), + rel_type=rel["rel_type"], + properties=rel["props"], + created_at=( + datetime.fromisoformat(rel["props"].get("created_at")) + if rel["props"].get("created_at") + else None + ), + ) + ) + + return RelationshipListResponse( + relationships=relationships, + total=total, + limit=params.limit, + offset=params.offset, + ) + + +def neo4j_update_relationship( + relationship_id: str, params: RelationshipUpdate +) -> RelationshipResponse: + """ + Update a relationship's properties. + + Authority: CanonKeeper only + Use Case: DL-14 + + Args: + relationship_id: Neo4j relationship ID + params: Update parameters + + Returns: + RelationshipResponse with updated data + + Raises: + ValueError: If relationship not found + """ + client = get_neo4j_client() + + # Verify relationship exists + existing = neo4j_get_relationship(relationship_id) + if not existing: + raise ValueError(f"Relationship {relationship_id} not found") + + # Update properties (preserve created_at) + updated_props = { + **params.properties, + "created_at": existing.created_at.isoformat() if existing.created_at else None, + } + + update_query = """ + MATCH ()-[r]->() + WHERE id(r) = $rel_id + SET r = $props + RETURN id(r) as rel_id + """ + + result = client.execute_write( + update_query, {"rel_id": int(relationship_id), "props": updated_props} + ) + + if not result: + raise ValueError(f"Failed to update relationship {relationship_id}") + + # Return updated relationship + updated = neo4j_get_relationship(relationship_id) + if not updated: + raise ValueError(f"Relationship {relationship_id} not found after update") + return updated + + +def neo4j_delete_relationship(relationship_id: str) -> Dict[str, Any]: + """ + Delete a relationship. + + Authority: CanonKeeper only + Use Case: DL-14 + + Args: + relationship_id: Neo4j relationship ID + + Returns: + Dict with deletion status + + Raises: + ValueError: If relationship not found + """ + client = get_neo4j_client() + + # Verify relationship exists + existing = neo4j_get_relationship(relationship_id) + if not existing: + raise ValueError(f"Relationship {relationship_id} not found") + + delete_query = """ + MATCH ()-[r]->() + WHERE id(r) = $rel_id + DELETE r + RETURN count(r) as deleted_count + """ + + result = client.execute_write(delete_query, {"rel_id": int(relationship_id)}) + + return { + "deleted": True, + "relationship_id": relationship_id, + "deleted_count": result[0]["deleted_count"] if result else 0, + } + + +# ============================================================================= +# STATE TAG TOOLS (DL-14) +# ============================================================================= + + +def neo4j_update_state_tags(params: StateTagUpdate) -> StateTagResponse: + """ + Update state tags on an entity instance atomically. + + Authority: CanonKeeper only + Use Case: DL-14 + + Args: + params: State tag update parameters + + Returns: + StateTagResponse with updated tags + + Raises: + ValueError: If entity not found or is an archetype + """ + client = get_neo4j_client() + + # Validate entity exists and is an instance + entity_check = client.execute_read( + """ + MATCH (e:Entity {id: $entity_id}) + RETURN e.id as id, e.entity_type as type + """, + {"entity_id": str(params.entity_id)}, + ) + + if not entity_check: + raise ValueError(f"Entity {params.entity_id} not found") + + if entity_check[0]["type"] == "archetype": + raise ValueError( + f"Cannot set state tags on archetype {params.entity_id}. " + "State tags are only valid on entity instances." + ) + + # Convert tags to strings + add_tag_strs = [tag.value for tag in params.add_tags] + remove_tag_strs = [tag.value for tag in params.remove_tags] + + # Update tags atomically + update_query = """ + MATCH (e:Entity {id: $entity_id}) + SET e.state_tags = + CASE + WHEN e.state_tags IS NULL THEN $add_tags + ELSE [tag IN coalesce(e.state_tags, []) + $add_tags WHERE NOT tag IN $remove_tags] + END + RETURN e.state_tags as tags + """ + + result = client.execute_write( + update_query, + { + "entity_id": str(params.entity_id), + "add_tags": add_tag_strs, + "remove_tags": remove_tag_strs, + }, + ) + + tags = result[0]["tags"] if result and result[0]["tags"] else [] + + return StateTagResponse(entity_id=params.entity_id, state_tags=tags) + + +def neo4j_get_state_tags(entity_id: UUID) -> StateTagResponse: + """ + Get current state tags for an entity. + + Authority: All agents + Use Case: DL-14 + + Args: + entity_id: Entity UUID + + Returns: + StateTagResponse with current tags + + Raises: + ValueError: If entity not found + """ + client = get_neo4j_client() + + query = """ + MATCH (e:Entity {id: $entity_id}) + RETURN e.state_tags as tags + """ + + result = client.execute_read(query, {"entity_id": str(entity_id)}) + + if not result: + raise ValueError(f"Entity {entity_id} not found") + + tags = result[0]["tags"] if result[0]["tags"] else [] + + return StateTagResponse(entity_id=entity_id, state_tags=tags) diff --git a/packages/data-layer/tests/test_tools/test_relationship_tools.py b/packages/data-layer/tests/test_tools/test_relationship_tools.py new file mode 100644 index 00000000..f067d977 --- /dev/null +++ b/packages/data-layer/tests/test_tools/test_relationship_tools.py @@ -0,0 +1,690 @@ +""" +Tests for Neo4j Relationship and State Tag operations (DL-14). + +Tests cover: +- neo4j_create_relationship +- neo4j_get_relationship +- neo4j_list_relationships +- neo4j_update_relationship +- neo4j_delete_relationship +- neo4j_update_state_tags +- neo4j_get_state_tags +""" + +from datetime import datetime, timezone +from unittest.mock import Mock, patch +from uuid import uuid4 + +import pytest + +from monitor_data.schemas.relationships import ( + RelationshipType, + Direction, + StateTag, + RelationshipCreate, + RelationshipUpdate, + RelationshipFilter, +) +from monitor_data.tools.neo4j_tools import ( + neo4j_create_relationship, + neo4j_get_relationship, + neo4j_list_relationships, + neo4j_update_relationship, + neo4j_delete_relationship, + neo4j_update_state_tags, + neo4j_get_state_tags, +) +from monitor_data.schemas.relationships import StateTagUpdate + + +# ============================================================================= +# TESTS: neo4j_create_relationship +# ============================================================================= + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_create_relationship_success(mock_get_client: Mock): + """Test creating a relationship between two entities.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + rel_id = "123" + + # Mock entity validation (both entities exist) + mock_client.execute_read.side_effect = [ + [{"id": str(entity1_id)}], # from_entity exists + [{"id": str(entity2_id)}], # to_entity exists + ] + + # Mock relationship creation + mock_client.execute_write.return_value = [ + { + "rel_id": rel_id, + "rel_type": "KNOWS", + "props": { + "since": "2020-01-01", + "strength": 5, + "created_at": datetime.now(timezone.utc).isoformat(), + }, + } + ] + + params = RelationshipCreate( + from_entity_id=entity1_id, + to_entity_id=entity2_id, + rel_type=RelationshipType.KNOWS, + properties={"since": "2020-01-01", "strength": 5}, + ) + + result = neo4j_create_relationship(params) + + assert result.from_entity_id == entity1_id + assert result.to_entity_id == entity2_id + assert result.rel_type == RelationshipType.KNOWS + assert result.properties["since"] == "2020-01-01" + assert result.properties["strength"] == 5 + assert result.relationship_id == rel_id + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_create_relationship_from_entity_not_found(mock_get_client: Mock): + """Test creating relationship with non-existent from_entity fails.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity2_id = uuid4() + + # Mock entity validation (from_entity doesn't exist) + mock_client.execute_read.return_value = [] + + params = RelationshipCreate( + from_entity_id=uuid4(), # Non-existent + to_entity_id=entity2_id, + rel_type=RelationshipType.KNOWS, + ) + + with pytest.raises(ValueError, match="From entity .* not found"): + neo4j_create_relationship(params) + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_create_relationship_to_entity_not_found(mock_get_client: Mock): + """Test creating relationship with non-existent to_entity fails.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + + # Mock entity validation (from exists, to doesn't) + mock_client.execute_read.side_effect = [ + [{"id": str(entity1_id)}], # from_entity exists + [], # to_entity doesn't exist + ] + + params = RelationshipCreate( + from_entity_id=entity1_id, + to_entity_id=uuid4(), # Non-existent + rel_type=RelationshipType.KNOWS, + ) + + with pytest.raises(ValueError, match="To entity .* not found"): + neo4j_create_relationship(params) + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_create_relationship_all_types(mock_get_client: Mock): + """Test creating relationships of all supported types.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + + # Test each relationship type + for rel_type in RelationshipType: + # Mock entity validation + mock_client.execute_read.side_effect = [ + [{"id": str(entity1_id)}], + [{"id": str(entity2_id)}], + ] + + # Mock relationship creation + mock_client.execute_write.return_value = [ + { + "rel_id": "123", + "rel_type": rel_type.value, + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + } + ] + + params = RelationshipCreate( + from_entity_id=entity1_id, + to_entity_id=entity2_id, + rel_type=rel_type, + ) + + result = neo4j_create_relationship(params) + assert result.rel_type == rel_type + + +# ============================================================================= +# TESTS: neo4j_get_relationship +# ============================================================================= + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_get_relationship_success(mock_get_client: Mock): + """Test retrieving a relationship by ID.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + rel_id = "123" + + mock_client.execute_read.return_value = [ + { + "rel_id": rel_id, + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "KNOWS", + "props": { + "since": "2020-01-01", + "strength": 5, + "created_at": datetime.now(timezone.utc).isoformat(), + }, + } + ] + + result = neo4j_get_relationship(rel_id) + + assert result is not None + assert result.relationship_id == rel_id + assert result.from_entity_id == entity1_id + assert result.to_entity_id == entity2_id + assert result.rel_type == RelationshipType.KNOWS + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_get_relationship_not_found(mock_get_client: Mock): + """Test retrieving non-existent relationship returns None.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + mock_client.execute_read.return_value = [] + + result = neo4j_get_relationship("999") + + assert result is None + + +# ============================================================================= +# TESTS: neo4j_list_relationships +# ============================================================================= + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_list_relationships_all(mock_get_client: Mock): + """Test listing all relationships without filters.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + + # Mock count query + mock_client.execute_read.side_effect = [ + [{"total": 2}], # count query + [ # data query + { + "rel_id": "1", + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "KNOWS", + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + }, + { + "rel_id": "2", + "from_id": str(entity2_id), + "to_id": str(entity1_id), + "rel_type": "ALLIED_WITH", + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + }, + ], + ] + + params = RelationshipFilter(limit=50) + result = neo4j_list_relationships(params) + + assert result.total == 2 + assert len(result.relationships) == 2 + assert result.relationships[0].rel_type == RelationshipType.KNOWS + assert result.relationships[1].rel_type == RelationshipType.ALLIED_WITH + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_list_relationships_by_entity_outgoing(mock_get_client: Mock): + """Test listing outgoing relationships from an entity.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + + # Mock count and data queries + mock_client.execute_read.side_effect = [ + [{"total": 1}], # count query + [ # data query + { + "rel_id": "1", + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "KNOWS", + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + } + ], + ] + + params = RelationshipFilter(entity_id=entity1_id, direction=Direction.OUTGOING) + result = neo4j_list_relationships(params) + + assert result.total == 1 + assert result.relationships[0].from_entity_id == entity1_id + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_list_relationships_by_entity_incoming(mock_get_client: Mock): + """Test listing incoming relationships to an entity.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + + # Mock count and data queries + mock_client.execute_read.side_effect = [ + [{"total": 1}], # count query + [ # data query + { + "rel_id": "1", + "from_id": str(entity2_id), + "to_id": str(entity1_id), + "rel_type": "KNOWS", + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + } + ], + ] + + params = RelationshipFilter(entity_id=entity1_id, direction=Direction.INCOMING) + result = neo4j_list_relationships(params) + + assert result.total == 1 + assert result.relationships[0].to_entity_id == entity1_id + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_list_relationships_by_type(mock_get_client: Mock): + """Test listing relationships filtered by type.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + + # Mock count and data queries + mock_client.execute_read.side_effect = [ + [{"total": 1}], # count query + [ # data query + { + "rel_id": "1", + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "ALLIED_WITH", + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + } + ], + ] + + params = RelationshipFilter(rel_type=RelationshipType.ALLIED_WITH) + result = neo4j_list_relationships(params) + + assert result.total == 1 + assert result.relationships[0].rel_type == RelationshipType.ALLIED_WITH + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_list_relationships_direction_both(mock_get_client: Mock): + """Test listing relationships in both directions.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + + # Mock count and data queries + mock_client.execute_read.side_effect = [ + [{"total": 2}], # count query + [ # data query + { + "rel_id": "1", + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "KNOWS", + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + }, + { + "rel_id": "2", + "from_id": str(entity2_id), + "to_id": str(entity1_id), + "rel_type": "KNOWS", + "props": {"created_at": datetime.now(timezone.utc).isoformat()}, + }, + ], + ] + + params = RelationshipFilter(entity_id=entity1_id, direction=Direction.BOTH) + result = neo4j_list_relationships(params) + + assert result.total == 2 + + +# ============================================================================= +# TESTS: neo4j_update_relationship +# ============================================================================= + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_update_relationship_properties(mock_get_client: Mock): + """Test updating relationship properties.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity1_id = uuid4() + entity2_id = uuid4() + rel_id = "123" + + # Mock get (to verify exists), then write, then get (to return updated) + created_at_iso = datetime.now(timezone.utc).isoformat() + mock_client.execute_read.side_effect = [ + # First get (verify exists) + [ + { + "rel_id": rel_id, + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "KNOWS", + "props": {"created_at": created_at_iso}, + } + ], + # Second get (return updated) + [ + { + "rel_id": rel_id, + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "KNOWS", + "props": { + "strength": 8, + "notes": "Updated relationship", + "created_at": created_at_iso, + }, + } + ], + ] + + # Mock write + mock_client.execute_write.return_value = [{"rel_id": rel_id}] + + params = RelationshipUpdate( + properties={"strength": 8, "notes": "Updated relationship"} + ) + + result = neo4j_update_relationship(rel_id, params) + + assert result.relationship_id == rel_id + assert result.properties["strength"] == 8 + assert result.properties["notes"] == "Updated relationship" + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_update_relationship_not_found(mock_get_client: Mock): + """Test updating non-existent relationship fails.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + # Mock get (not found) + mock_client.execute_read.return_value = [] + + params = RelationshipUpdate(properties={"strength": 8}) + + with pytest.raises(ValueError, match="Relationship .* not found"): + neo4j_update_relationship("999", params) + + +# ============================================================================= +# TESTS: neo4j_delete_relationship +# ============================================================================= + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_delete_relationship_success(mock_get_client: Mock): + """Test deleting a relationship.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + rel_id = "123" + entity1_id = uuid4() + entity2_id = uuid4() + + # Mock get (verify exists) + mock_client.execute_read.return_value = [ + { + "rel_id": rel_id, + "from_id": str(entity1_id), + "to_id": str(entity2_id), + "rel_type": "KNOWS", + "props": {}, + } + ] + + # Mock successful delete + mock_client.execute_write.return_value = [{"deleted_count": 1}] + + result = neo4j_delete_relationship(rel_id) + + assert result["deleted"] is True + assert result["deleted_count"] == 1 + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_delete_relationship_not_found(mock_get_client: Mock): + """Test deleting non-existent relationship fails.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + # Mock get (not found) + mock_client.execute_read.return_value = [] + + with pytest.raises(ValueError, match="Relationship .* not found"): + neo4j_delete_relationship("999") + + +# ============================================================================= +# TESTS: neo4j_update_state_tags +# ============================================================================= + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_update_state_tags_add(mock_get_client: Mock): + """Test adding state tags to an entity.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity_id = uuid4() + + # Mock entity validation (is instance) + mock_client.execute_read.return_value = [{"id": str(entity_id), "type": "instance"}] + + # Mock tag update + mock_client.execute_write.return_value = [{"tags": ["alive", "wounded"]}] + + params = StateTagUpdate( + entity_id=entity_id, + add_tags=[StateTag.ALIVE, StateTag.WOUNDED], + ) + + result = neo4j_update_state_tags(params) + + assert result.entity_id == entity_id + assert StateTag.ALIVE in result.state_tags + assert StateTag.WOUNDED in result.state_tags + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_update_state_tags_remove(mock_get_client: Mock): + """Test removing state tags from an entity.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity_id = uuid4() + + # Mock entity validation (is instance with existing tags) + mock_client.execute_read.return_value = [{"id": str(entity_id), "type": "instance"}] + + # Mock tag update + mock_client.execute_write.return_value = [{"tags": ["alive"]}] + + params = StateTagUpdate( + entity_id=entity_id, + remove_tags=[StateTag.WOUNDED], + ) + + result = neo4j_update_state_tags(params) + + assert result.entity_id == entity_id + assert StateTag.ALIVE in result.state_tags + assert StateTag.WOUNDED not in result.state_tags + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_update_state_tags_add_and_remove(mock_get_client: Mock): + """Test adding and removing state tags atomically.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity_id = uuid4() + + # Mock entity validation + mock_client.execute_read.return_value = [{"id": str(entity_id), "type": "instance"}] + + # Mock tag update + mock_client.execute_write.return_value = [ + {"tags": ["dead"]} # removed wounded, added dead + ] + + params = StateTagUpdate( + entity_id=entity_id, + add_tags=[StateTag.DEAD], + remove_tags=[StateTag.WOUNDED], + ) + + result = neo4j_update_state_tags(params) + + assert result.entity_id == entity_id + assert StateTag.DEAD in result.state_tags + assert StateTag.WOUNDED not in result.state_tags + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_update_state_tags_archetype_fails(mock_get_client: Mock): + """Test that state tags cannot be applied to archetypes.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity_id = uuid4() + + # Mock entity validation (is archetype) + mock_client.execute_read.return_value = [ + {"id": str(entity_id), "type": "archetype"} + ] + + params = StateTagUpdate( + entity_id=entity_id, + add_tags=[StateTag.ALIVE], + ) + + with pytest.raises(ValueError, match="Cannot set state tags on archetype"): + neo4j_update_state_tags(params) + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_update_state_tags_entity_not_found(mock_get_client: Mock): + """Test updating state tags for non-existent entity fails.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + # Mock entity not found + mock_client.execute_read.return_value = [] + + params = StateTagUpdate( + entity_id=uuid4(), + add_tags=[StateTag.ALIVE], + ) + + with pytest.raises(ValueError, match="Entity .* not found"): + neo4j_update_state_tags(params) + + +# ============================================================================= +# TESTS: neo4j_get_state_tags +# ============================================================================= + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_get_state_tags_success(mock_get_client: Mock): + """Test retrieving state tags for an entity.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity_id = uuid4() + + mock_client.execute_read.return_value = [{"tags": ["alive", "wounded", "prone"]}] + + result = neo4j_get_state_tags(entity_id) + + assert result.entity_id == entity_id + assert len(result.state_tags) == 3 + assert StateTag.ALIVE in result.state_tags + assert StateTag.WOUNDED in result.state_tags + assert StateTag.PRONE in result.state_tags + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_get_state_tags_empty(mock_get_client: Mock): + """Test retrieving state tags for entity with no tags.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + entity_id = uuid4() + + mock_client.execute_read.return_value = [{"tags": []}] + + result = neo4j_get_state_tags(entity_id) + + assert result.entity_id == entity_id + assert len(result.state_tags) == 0 + + +@patch("monitor_data.tools.neo4j_tools.get_neo4j_client") +def test_get_state_tags_entity_not_found(mock_get_client: Mock): + """Test retrieving state tags for non-existent entity fails.""" + mock_client = Mock() + mock_get_client.return_value = mock_client + + # Mock entity not found + mock_client.execute_read.return_value = [] + + with pytest.raises(ValueError, match="Entity .* not found"): + neo4j_get_state_tags(uuid4()) From a37bc9cfeb3f52939cbe793385c7e9e7cc9c1915 Mon Sep 17 00:00:00 2001 From: spuentesp Date: Mon, 5 Jan 2026 10:23:05 -0300 Subject: [PATCH 2/2] fix(data-layer): Address all PR review comments for DL-14 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses all 14 review comments from Copilot code review: ## Critical Fixes **P1: Archetype validation bug** (line 3888) - Fixed query to check `e.is_archetype` instead of `e.entity_type` - Now correctly prevents state tags on archetypes - Updated test mocks to use `{"is_archetype": True/False}` **Cypher injection prevention** (line 3706) - Parameterized rel_type filter in list_relationships - Changed from string interpolation to query parameter - Note: CREATE statement (line 3602) cannot parameterize rel_type due to Neo4j limitation, but enum validation provides safety ## Data Integrity Fixes **DELETE count query** (line 3838) - Fixed query to use `WITH r DELETE r RETURN count(*)` - Previously returned 0 because r was deleted before count - Now correctly returns deleted count **Duplicate tags handling** (line 3901) - Rewrote state_tags update query to use REDUCE for deduplication - Remove tags first, then add, then deduplicate - If same tag in both add/remove, addition takes precedence ## Validation Improvements **Relationship ID validation** (lines 3651, 3798, 3842) - Added try-except blocks for int conversion - Provides clear error: "Invalid relationship ID format: must be a numeric string" - Applied to get, update, and delete functions **Empty tags validation** (line 3917) - Requires at least one of add_tags or remove_tags to be non-empty - Prevents no-op database write operations **Self-reference validation** (line 3627) - Prevents self-referencing relationships for KNOWS, ALLIED_WITH, HOSTILE_TO - Allows OWNS and other types where self-reference may be valid - Provides clear error message ## Architecture Notes - CREATE statement line 3602: Cannot parameterize relationship type due to Neo4j Cypher limitations. Enum validation provides security. - Race condition comments (lines 3593, 3808, 3841): Acknowledged but acceptable given entity/relationship validation and Neo4j ACID properties All 23 tests passing ✅ 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Sonnet 4.5 --- .../src/monitor_data/tools/neo4j_tools.py | 73 +++++++++++++++---- .../test_tools/test_relationship_tools.py | 14 +++- 2 files changed, 70 insertions(+), 17 deletions(-) diff --git a/packages/data-layer/src/monitor_data/tools/neo4j_tools.py b/packages/data-layer/src/monitor_data/tools/neo4j_tools.py index 80736c85..9649c2df 100644 --- a/packages/data-layer/src/monitor_data/tools/neo4j_tools.py +++ b/packages/data-layer/src/monitor_data/tools/neo4j_tools.py @@ -66,6 +66,7 @@ SetActivePC, ) from monitor_data.schemas.relationships import ( + RelationshipType, RelationshipCreate, RelationshipUpdate, RelationshipResponse, @@ -3592,6 +3593,18 @@ def neo4j_create_relationship(params: RelationshipCreate) -> RelationshipRespons if not to_exists: raise ValueError(f"To entity {params.to_entity_id} not found") + # Validate no self-reference for relationship types where it doesn't make sense + if params.from_entity_id == params.to_entity_id: + # OWNS might be valid (e.g., recursive ownership), but most types are not + if params.rel_type in ( + RelationshipType.KNOWS, + RelationshipType.ALLIED_WITH, + RelationshipType.HOSTILE_TO, + ): + raise ValueError( + f"Self-referencing relationships are not allowed for {params.rel_type.value}" + ) + # Create relationship with properties now = datetime.now(timezone.utc) props = {**params.properties, "created_at": now.isoformat()} @@ -3648,7 +3661,14 @@ def neo4j_get_relationship(relationship_id: str) -> Optional[RelationshipRespons type(r) as rel_type, properties(r) as props """ - result = client.execute_read(query, {"rel_id": int(relationship_id)}) + try: + rel_id_int = int(relationship_id) + except (TypeError, ValueError): + raise ValueError( + "Invalid relationship ID format: must be a numeric string" + ) from None + + result = client.execute_read(query, {"rel_id": rel_id_int}) if not result: return None @@ -3703,7 +3723,8 @@ def neo4j_list_relationships( query_params["entity_id"] = str(params.entity_id) if params.rel_type: - where_clauses.append(f"type(r) = '{params.rel_type.value}'") + where_clauses.append("type(r) = $rel_type") + query_params["rel_type"] = params.rel_type.value where_clause = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else "" @@ -3794,8 +3815,15 @@ def neo4j_update_relationship( RETURN id(r) as rel_id """ + try: + rel_id_int = int(relationship_id) + except (TypeError, ValueError): + raise ValueError( + "Invalid relationship ID format: must be a numeric string" + ) from None + result = client.execute_write( - update_query, {"rel_id": int(relationship_id), "props": updated_props} + update_query, {"rel_id": rel_id_int, "props": updated_props} ) if not result: @@ -3834,11 +3862,19 @@ def neo4j_delete_relationship(relationship_id: str) -> Dict[str, Any]: delete_query = """ MATCH ()-[r]->() WHERE id(r) = $rel_id + WITH r DELETE r - RETURN count(r) as deleted_count + RETURN count(*) as deleted_count """ - result = client.execute_write(delete_query, {"rel_id": int(relationship_id)}) + try: + rel_id_int = int(relationship_id) + except (TypeError, ValueError): + raise ValueError( + "Invalid relationship ID format: must be a numeric string" + ) from None + + result = client.execute_write(delete_query, {"rel_id": rel_id_int}) return { "deleted": True, @@ -3874,7 +3910,7 @@ def neo4j_update_state_tags(params: StateTagUpdate) -> StateTagResponse: entity_check = client.execute_read( """ MATCH (e:Entity {id: $entity_id}) - RETURN e.id as id, e.entity_type as type + RETURN e.id as id, e.is_archetype as is_archetype """, {"entity_id": str(params.entity_id)}, ) @@ -3882,24 +3918,35 @@ def neo4j_update_state_tags(params: StateTagUpdate) -> StateTagResponse: if not entity_check: raise ValueError(f"Entity {params.entity_id} not found") - if entity_check[0]["type"] == "archetype": + if entity_check[0]["is_archetype"]: raise ValueError( f"Cannot set state tags on archetype {params.entity_id}. " "State tags are only valid on entity instances." ) + # Validate at least one operation + if not params.add_tags and not params.remove_tags: + raise ValueError("At least one of add_tags or remove_tags must be non-empty") + # Convert tags to strings add_tag_strs = [tag.value for tag in params.add_tags] remove_tag_strs = [tag.value for tag in params.remove_tags] - # Update tags atomically + # Update tags atomically (remove first, then add, then deduplicate) + # If same tag in both add and remove, addition takes precedence update_query = """ MATCH (e:Entity {id: $entity_id}) - SET e.state_tags = - CASE - WHEN e.state_tags IS NULL THEN $add_tags - ELSE [tag IN coalesce(e.state_tags, []) + $add_tags WHERE NOT tag IN $remove_tags] - END + WITH e, + [tag IN coalesce(e.state_tags, []) WHERE NOT tag IN $remove_tags] as after_remove + SET e.state_tags = + REDUCE( + s = [], + t IN (after_remove + $add_tags) | + CASE + WHEN t IN s THEN s + ELSE s + t + END + ) RETURN e.state_tags as tags """ diff --git a/packages/data-layer/tests/test_tools/test_relationship_tools.py b/packages/data-layer/tests/test_tools/test_relationship_tools.py index f067d977..921bc28e 100644 --- a/packages/data-layer/tests/test_tools/test_relationship_tools.py +++ b/packages/data-layer/tests/test_tools/test_relationship_tools.py @@ -525,7 +525,9 @@ def test_update_state_tags_add(mock_get_client: Mock): entity_id = uuid4() # Mock entity validation (is instance) - mock_client.execute_read.return_value = [{"id": str(entity_id), "type": "instance"}] + mock_client.execute_read.return_value = [ + {"id": str(entity_id), "is_archetype": False} + ] # Mock tag update mock_client.execute_write.return_value = [{"tags": ["alive", "wounded"]}] @@ -551,7 +553,9 @@ def test_update_state_tags_remove(mock_get_client: Mock): entity_id = uuid4() # Mock entity validation (is instance with existing tags) - mock_client.execute_read.return_value = [{"id": str(entity_id), "type": "instance"}] + mock_client.execute_read.return_value = [ + {"id": str(entity_id), "is_archetype": False} + ] # Mock tag update mock_client.execute_write.return_value = [{"tags": ["alive"]}] @@ -577,7 +581,9 @@ def test_update_state_tags_add_and_remove(mock_get_client: Mock): entity_id = uuid4() # Mock entity validation - mock_client.execute_read.return_value = [{"id": str(entity_id), "type": "instance"}] + mock_client.execute_read.return_value = [ + {"id": str(entity_id), "is_archetype": False} + ] # Mock tag update mock_client.execute_write.return_value = [ @@ -607,7 +613,7 @@ def test_update_state_tags_archetype_fails(mock_get_client: Mock): # Mock entity validation (is archetype) mock_client.execute_read.return_value = [ - {"id": str(entity_id), "type": "archetype"} + {"id": str(entity_id), "is_archetype": True} ] params = StateTagUpdate(