agent-management/backend/src/test/java/com/agent/management/rag/structured/VannaSqlGenerationServiceTest.java

package com.agent.management.rag.structured;

import com.agent.management.model.entity.DataSource;
import com.agent.management.rag.bridge.RagAiBridgeClient;
import com.agent.management.rag.capability.*;
import com.agent.management.rag.memory.RagExampleMemoryService;
import com.agent.management.rag.model.RagQuery;
import com.agent.management.rag.model.RagSourceType;
import com.agent.management.service.DataSourceService;
import com.agent.management.service.SchemaExplorerService;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;

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

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

class VannaSqlGenerationServiceTest {
    @Test
    void sendsOnlyWhitelistedTableDdlAndConfiguredContext() {
        RagAiBridgeClient bridge = mock(RagAiBridgeClient.class);
        DataSourceService dataSources = mock(DataSourceService.class);
        RagCapabilityProfileService profiles = mock(RagCapabilityProfileService.class);
        RagSchemaLinker linker = mock(RagSchemaLinker.class);
        RagExampleMemoryService examples = mock(RagExampleMemoryService.class);
        DataSource source = new DataSource();
        source.setType("MYSQL");
        when(dataSources.getDecrypted(1L)).thenReturn(source);
        SqlSchemaSnapshot snapshot = new SqlSchemaSnapshot(List.of(
                new SqlSchemaSnapshot.TableSchema("rescue_forces", "", List.of(
                        new SqlSchemaSnapshot.ColumnSchema("name", "VARCHAR", "力量名称"),
                        new SqlSchemaSnapshot.ColumnSchema("status", "VARCHAR", "状态"))),
                new SqlSchemaSnapshot.TableSchema("secret_table", "", List.of(
                        new SqlSchemaSnapshot.ColumnSchema("secret", "VARCHAR", "")))
        ));
        when(profiles.get(RagSourceType.STRUCTURED_DATA, "1"))
                .thenReturn(new RagCapabilityProfile(RagSourceType.STRUCTURED_DATA,"1","v",java.time.Instant.now(),snapshot,null,List.of()));
        when(linker.selectSqlTables(eq("可用力量"), eq(snapshot), anyList(), anyMap())).thenReturn(List.of("rescue_forces"));
        when(examples.findSimilar(any(), anyString(), anyString(), anyInt())).thenReturn(List.of());
        when(bridge.textToSql(any())).thenReturn(Map.of("sql", "SELECT name FROM rescue_forces"));
        RagQuery query = new RagQuery();
        query.setQuery("可用力量");
        query.setFilters(Map.of(
                "allowTextToSql", true,
                "tableWhitelist", List.of("rescue_forces"),
                "documentation", "READY 表示可用",
                "examples", List.of(Map.of("question", "q", "sql", "SELECT 1")),
                "maxRows", 20
        ));

        assertThat(new VannaSqlGenerationService(bridge, dataSources, profiles, linker, mock(RagSemanticCatalogService.class), examples, new RagEntityMentionExtractor()).generateSql(query, 1L))
                .contains("SELECT name FROM rescue_forces");

        ArgumentCaptor<Map<String, Object>> request = ArgumentCaptor.forClass(Map.class);
        verify(bridge).textToSql(request.capture());
        assertThat(String.valueOf(request.getValue().get("ddl")))
                .contains("rescue_forces", "name VARCHAR", "status VARCHAR")
                .doesNotContain("secret_table");
        assertThat(request.getValue()).containsEntry("documentation", "READY 表示可用")
                .containsEntry("maxRows", 20);
    }
}