package com.agent.management.rag.capability;

import com.agent.management.rag.graph.GraphSchemaSnapshot;
import com.agent.management.rag.model.RagSourceType;
import org.junit.jupiter.api.Test;

import java.time.Instant;
import java.util.List;
import java.util.Map;

import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.*;

class GraphBusinessSubgraphSelectorTest {
    @Test
    void expandsSemanticSeedsThroughRelationshipEndpoints() {
        RagSemanticCatalogService semantic = mock(RagSemanticCatalogService.class);
        RagCapabilityProfile profile = new RagCapabilityProfile(RagSourceType.GRAPH, "1", "v", Instant.now(), null, schema(), List.of());
        when(semantic.rank("任务风险", profile, "label:", 8)).thenReturn(List.of("Mission"));
        when(semantic.rank("任务风险", profile, "relationship:", 8)).thenReturn(List.of("THREATENS"));

        GraphSchemaSnapshot selected = new GraphBusinessSubgraphSelector(semantic).select("任务风险", profile, schema(), 8, 8);

        assertThat(selected.labelNames()).contains("Mission", "RiskEvent").doesNotContain("Report");
        assertThat(selected.relationshipTypeNames()).containsExactly("THREATENS");
    }

    private static GraphSchemaSnapshot schema() {
        return new GraphSchemaSnapshot(List.of(
                new GraphSchemaSnapshot.NodeSchema("Mission", Map.of("name", List.of("String"))),
                new GraphSchemaSnapshot.NodeSchema("RiskEvent", Map.of("name", List.of("String"))),
                new GraphSchemaSnapshot.NodeSchema("Report", Map.of("title", List.of("String")))),
                List.of(new GraphSchemaSnapshot.RelationshipSchema("THREATENS", List.of("RiskEvent"), List.of("Mission"), Map.of())));
    }
}
