agent-management/backend/src/test/java/com/agent/management/rag/graph/GraphSchemaSnapshotTest.java

package com.agent.management.rag.graph;

import org.junit.jupiter.api.Test;

import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;

class GraphSchemaSnapshotTest {
    @Test
    void whitelistIntersectsRealSchemaAndKeepsDirectionAndProperties() {
        GraphSchemaSnapshot schema = new GraphSchemaSnapshot(
                List.of(
                        new GraphSchemaSnapshot.NodeSchema("Mission", Map.of("name", List.of("STRING"))),
                        new GraphSchemaSnapshot.NodeSchema("Stage", Map.of("name", List.of("STRING"), "sequence", List.of("INTEGER"), "goal", List.of("STRING"))),
                        new GraphSchemaSnapshot.NodeSchema("Secret", Map.of("value", List.of("STRING")))
                ),
                List.of(
                        new GraphSchemaSnapshot.RelationshipSchema("HAS_STAGE", List.of("Mission"), List.of("Stage"), Map.of()),
                        new GraphSchemaSnapshot.RelationshipSchema("SECRET_LINK", List.of("Mission"), List.of("Secret"), Map.of())
                )
        );

        GraphSchemaSnapshot filtered = schema.filter(
                List.of("Mission", "Stage"),
                List.of("HAS_STAGE"),
                Map.of("Mission", List.of("name"), "Stage", List.of("name", "sequence", "goal"))
        );

        assertThat(filtered.labelNames()).containsExactly("Mission", "Stage");
        assertThat(filtered.relationshipTypeNames()).containsExactly("HAS_STAGE");
        assertThat(filtered.node("Stage").orElseThrow().properties()).containsOnlyKeys("name", "sequence", "goal");
        assertThat(filtered.relationships()).singleElement().satisfies(rel -> {
            assertThat(rel.startLabels()).containsExactly("Mission");
            assertThat(rel.endLabels()).containsExactly("Stage");
        });
    }
}