content: /**
   * Copyright (C) 2024 AIDC-AI
   *
   * Licensed under the Apache License, Version 2.0 (the "License");
   * you may not use this file except in compliance with the License.
   * You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  package com.alibaba.langengine.tinkerpop.vectorstore;
  
  import com.alibaba.langengine.core.embeddings.Embeddings;
  import com.alibaba.langengine.core.indexes.Document;
  import com.alibaba.langengine.tinkerpop.vectorstore.service.*;
  import org.junit.jupiter.api.Test;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.AfterEach;
  import org.junit.jupiter.api.Disabled;
  import org.junit.jupiter.api.extension.ExtendWith;
  import org.mockito.Mock;
  import org.mockito.junit.jupiter.MockitoExtension;
  
  import java.util.*;
  
  import static org.junit.jupiter.api.Assertions.*;
  import static org.mockito.Mockito.*;
  
  
  @ExtendWith(MockitoExtension.class)
  public class TinkerPopMockTest {
  
      @Mock
      private Embeddings mockEmbedding;
  
      @Mock
      private TinkerPopService mockService;
  
      @Mock
      private TinkerPopClient mockClient;
  
      private TinkerPop tinkerPop;
      private final String TEST_COLLECTION_ID = "test_collection";
      private final String TEST_SERVER_URL = "ws://test-server:8182";
  
      @BeforeEach
      public void setUp() {
          // 创建一个 TinkerPop 实例并设置 mock 服务
          tinkerPop = new TinkerPop(mockEmbedding, TEST_COLLECTION_ID);
  
          // 使用反射设置 mock 服务
          setPrivateField(tinkerPop, "_service", mockService);
          setPrivateField(tinkerPop, "_client", mockClient);
      }
  
      @AfterEach
      public void tearDown() {
          if (tinkerPop != null) {
              tinkerPop.close();
          }
      }
  
      @Test
      public void testConstructorWithNullCollectionId() {
          TinkerPop tinkerPopWithNullId = new TinkerPop(mockEmbedding, null);
          assertNotNull(tinkerPopWithNullId.getCollectionId());
          assertNotEquals("", tinkerPopWithNullId.getCollectionId());
          tinkerPopWithNullId.close();
      }
  
      @Test
      @Disabled("Requires a running TinkerPop server at test-server:8182")
      public void testConstructorWithCustomServerUrl() {
          TinkerPop customTinkerPop = new TinkerPop(TEST_SERVER_URL, mockEmbedding, TEST_COLLECTION_ID);
          assertNotNull(customTinkerPop);
          assertEquals(TEST_COLLECTION_ID, customTinkerPop.getCollectionId());
          assertEquals(mockEmbedding, customTinkerPop.getEmbedding());
          customTinkerPop.close();
      }
  
      @Test
      @Disabled("Requires a running TinkerPop server at test-server:8182")
      public void testConstructorWithTimeouts() {
          TinkerPop timeoutTinkerPop = new TinkerPop(
              TEST_SERVER_URL, mockEmbedding, TEST_COLLECTION_ID, 5000, 10000);
          assertNotNull(timeoutTinkerPop);
          assertEquals(TEST_COLLECTION_ID, timeoutTinkerPop.getCollectionId());
          timeoutTinkerPop.close();
      }
  
      @Test
      public void testAddDocumentsWithValidDocuments() {
          // 准备测试数据
          Document doc1 = new Document();
          doc1.setUniqueId("doc1");
          doc1.setPageContent("Test content 1");
          doc1.setEmbedding(Arrays.asList(0.1, 0.2, 0.3));
          doc1.setMetadata(Map.of("source", "test1.txt"));
  
          Document doc2 = new Document();
          doc2.setUniqueId("doc2");
          doc2.setPageContent("Test content 2");
          doc2.setEmbedding(Arrays.asList(0.4, 0.5, 0.6));
          doc2.setMetadata(Map.of("source", "test2.txt"));
  
          List<Document> documents = Arrays.asList(doc1, doc2);
  
          // Mock 服务行为
          doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));
  
          // 执行测试
          assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));
  
          // 验证服务被调用
          verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
      }
  
      @Test
      public void testAddDocumentsWithEmptyContent() {
          // 准备测试数据 - 包含空内容
          Document docWithEmptyContent = new Document();
          docWithEmptyContent.setUniqueId("empty-doc");
          docWithEmptyContent.setPageContent(""); // 空内容
  
          Document docWithNullContent = new Document();
          docWithNullContent.setUniqueId("null-doc");
          docWithNullContent.setPageContent(null); // null内容
  
          List<Document> documents = Arrays.asList(docWithEmptyContent, docWithNullContent);
  
          doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));
  
          assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));
          verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
      }
  
      @Test
      public void testAddDocumentsWithAutoGeneratedIds() {
          Document docWithoutId = new Document();
          docWithoutId.setPageContent("Content without ID");
  
          List<Document> documents = Arrays.asList(docWithoutId);
  
          doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));
  
          assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));
  
          // 验证ID被自动生成
          assertNotNull(docWithoutId.getUniqueId());
          assertFalse(docWithoutId.getUniqueId().isEmpty());
  
          verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
      }
  
      @Test
      public void testAddDocumentsWithEmbeddingGeneration() {
          Document docWithoutEmbedding = new Document();
          docWithoutEmbedding.setUniqueId("doc-no-embedding");
          docWithoutEmbedding.setPageContent("Content without embedding");
  
          List<Document> documents = Arrays.asList(docWithoutEmbedding);
  
          // Mock embedding 服务生成嵌入向量
          Document embeddedDoc = new Document();
          embeddedDoc.setEmbedding(Arrays.asList(0.1, 0.2, 0.3));
          when(mockEmbedding.embedTexts(anyList())).thenReturn(Arrays.asList(embeddedDoc));
  
          doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));
  
          assertDoesNotThrow(() -> tinkerPop.addDocuments(documents));
  
          verify(mockEmbedding).embedTexts(anyList());
          verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
      }
  
      @Test
      public void testAddTexts() {
          List<String> texts = Arrays.asList("Text 1", "Text 2", "Text 3");
          List<Map<String, Object>> metadatas = Arrays.asList(
              Map.of("source", "file1.txt"),
              Map.of("source", "file2.txt"),
              Map.of("source", "file3.txt")
          );
          List<String> ids = Arrays.asList("id1", "id2", "id3");
  
          // Mock embedding 服务
          List<Document> embeddedDocs = Arrays.asList(
              createMockDocument("id1", Arrays.asList(0.1, 0.2)),
              createMockDocument("id2", Arrays.asList(0.3, 0.4)),
              createMockDocument("id3", Arrays.asList(0.5, 0.6))
          );
          when(mockEmbedding.embedTexts(texts)).thenReturn(embeddedDocs);
  
          doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));
  
          List<String> resultIds = tinkerPop.addTexts(texts, metadatas, ids);
  
          assertEquals(ids, resultIds);
          verify(mockEmbedding).embedTexts(texts);
          verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
      }
  
      @Test
      public void testAddTextsWithNullIds() {
          List<String> texts = Arrays.asList("Text 1", "Text 2");
          List<Map<String, Object>> metadatas = Arrays.asList(
              Map.of("source", "file1.txt"),
              Map.of("source", "file2.txt")
          );
  
          // Mock embedding 服务
          List<Document> embeddedDocs = Arrays.asList(
              createMockDocument("auto-id-1", Arrays.asList(0.1, 0.2)),
              createMockDocument("auto-id-2", Arrays.asList(0.3, 0.4))
          );
          when(mockEmbedding.embedTexts(texts)).thenReturn(embeddedDocs);
  
          doNothing().when(mockService).addDocuments(any(TinkerPopAddRequest.class));
  
          List<String> resultIds = tinkerPop.addTexts(texts, metadatas, null);
  
          assertEquals(2, resultIds.size());
          // 验证ID被自动生成
          resultIds.forEach(id -> {
              assertNotNull(id);
              assertFalse(id.isEmpty());
          });
  
          verify(mockEmbedding).embedTexts(texts);
          verify(mockService).addDocuments(any(TinkerPopAddRequest.class));
      }
  
      @Test
      public void testSimilaritySearch() {
          String query = "test query";
          int k = 5;
  
          // Mock embedding 服务为查询生成嵌入向量
          Document queryEmbedDoc = createMockDocument("query", Arrays.asList(0.1, 0.2, 0.3));
          when(mockEmbedding.embedTexts(Arrays.asList(query))).thenReturn(Arrays.asList(queryEmbedDoc));
  
          // Mock 服务查询响应
          TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
          mockResponse.setIds(Arrays.asList("doc1", "doc2"));
          mockResponse.setTexts(Arrays.asList("result text 1", "result text 2"));
          mockResponse.setDistances(Arrays.asList(0.1, 0.2));
          mockResponse.setMetadatas(Arrays.asList(
              Map.of("source", "file1.txt"),
              Map.of("source", "file2.txt")
          ));
  
          when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);
  
          List<Document> results = tinkerPop.similaritySearch(query, k, null, null);
  
          assertEquals(2, results.size());
          assertEquals("doc1", results.get(0).getUniqueId());
          assertEquals("doc2", results.get(1).getUniqueId());
          assertEquals("result text 1", results.get(0).getPageContent());
          assertEquals("result text 2", results.get(1).getPageContent());
  
          verify(mockEmbedding).embedTexts(Arrays.asList(query));
          verify(mockService).queryDocuments(any(TinkerPopQueryRequest.class));
      }
  
      @Test
      public void testSimilaritySearchWithMaxDistance() {
          String query = "test query";
          int k = 5;
          Double maxDistance = 0.15;
  
          // Mock embedding 服务
          Document queryEmbedDoc = createMockDocument("query", Arrays.asList(0.1, 0.2, 0.3));
          when(mockEmbedding.embedTexts(Arrays.asList(query))).thenReturn(Arrays.asList(queryEmbedDoc));
  
          // Mock 服务查询响应 - 包含距离超过阈值的结果
          TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
          mockResponse.setIds(Arrays.asList("doc1", "doc2", "doc3"));
          mockResponse.setTexts(Arrays.asList("result 1", "result 2", "result 3"));
          mockResponse.setDistances(Arrays.asList(0.1, 0.2, 0.3)); // doc2 和 doc3 超过阈值
          mockResponse.setMetadatas(Arrays.asList(
              Map.of("source", "file1.txt"),
              Map.of("source", "file2.txt"),
              Map.of("source", "file3.txt")
          ));
  
          when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);
  
          List<Document> results = tinkerPop.similaritySearch(query, k, maxDistance, null);
  
          // 只有 doc1 的距离 <= 0.15
          assertEquals(1, results.size());
          assertEquals("doc1", results.get(0).getUniqueId());
          assertEquals(0.1, results.get(0).getScore());
  
          verify(mockEmbedding).embedTexts(Arrays.asList(query));
          verify(mockService).queryDocuments(any(TinkerPopQueryRequest.class));
      }
  
      @Test
      public void testSimilaritySearchWithNullEmbedding() {
          tinkerPop.setEmbedding(null); // 设置为 null
          String query = "test query";
          int k = 5;
  
          TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
          mockResponse.setIds(Arrays.asList("doc1"));
          mockResponse.setTexts(Arrays.asList("result text"));
          mockResponse.setDistances(Arrays.asList(0.1));
          mockResponse.setMetadatas(Arrays.asList(Map.of("source", "file.txt")));
  
          when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);
  
          List<Document> results = tinkerPop.similaritySearch(query, k, null, null);
  
          assertEquals(1, results.size());
          verify(mockService).queryDocuments(any(TinkerPopQueryRequest.class));
          // 验证没有调用 embedding 服务
          verifyNoInteractions(mockEmbedding);
      }
  
      @Test
      public void testSimilaritySearchWithHtmlFiltering() {
          String query = "test query";
  
          Document queryEmbedDoc = createMockDocument("query", Arrays.asList(0.1, 0.2, 0.3));
          when(mockEmbedding.embedTexts(Arrays.asList(query))).thenReturn(Arrays.asList(queryEmbedDoc));
  
          // Mock 响应包含 HTML 标签的文本
          TinkerPopQueryResponse mockResponse = new TinkerPopQueryResponse();
          mockResponse.setIds(Arrays.asList("doc1"));
          mockResponse.setTexts(Arrays.asList("<p>This is <b>bold</b> text &amp; escaped</p>"));
          mockResponse.setDistances(Arrays.asList(0.1));
          mockResponse.setMetadatas(Arrays.asList(Map.of("source", "file.txt")));
  
          when(mockService.queryDocuments(any(TinkerPopQueryRequest.class))).thenReturn(mockResponse);
  
          List<Document> results = tinkerPop.similaritySearch(query, 5, null, null);
  
          assertEquals(1, results.size());
          // 验证 HTML 标签被过滤，HTML 实体被解码
          assertEquals("This is bold text & escaped", results.get(0).getPageContent());
      }
  
      @Test
      public void testAddDocumentsThrowsException() {
          Document doc = new Document();
          doc.setUniqueId("doc1");
          doc.setPageContent("content");
          doc.setMetadata(Map.of("key", "value"));
  
          doThrow(new RuntimeException("Service error")).when(mockService).addDocuments(any(TinkerPopAddRequest.class));
  
          RuntimeException exception = assertThrows(RuntimeException.class, () -> {
              tinkerPop.addDocuments(Arrays.asList(doc));
          });
  
          assertTrue(exception.getMessage().contains("Service error"));
      }
  
      @Test
      public void testSimilaritySearchThrowsException() {
          when(mockService.queryDocuments(any(TinkerPopQueryRequest.class)))
              .thenThrow(new RuntimeException("Query error"));
  
          RuntimeException exception = assertThrows(RuntimeException.class, () -> {
              tinkerPop.similaritySearch("query", 5, null, null);
          });
  
          assertTrue(exception.getMessage().contains("Query error"));
      }
  
      @Test
      public void testClose() {
          doNothing().when(mockService).close();
  
          assertDoesNotThrow(() -> tinkerPop.close());
          verify(mockService).close();
      }
  
      @Test
      public void testFilterMethod() {
          // 使用反射测试私有的 filter 方法
          String htmlText = "<p>Test <b>bold</b> text &lt;script&gt; &amp; more</p>";
          String filtered = invokeFilterMethod(htmlText);
  
          assertEquals("Test bold text <script> & more", filtered);
      }
  
      @Test
      public void testFilterWithNullInput() {
          String filtered = invokeFilterMethod(null);
          assertEquals("", filtered);
      }
  
      // 辅助方法
      private Document createMockDocument(String id, List<Double> embedding) {
          Document doc = new Document();
          doc.setUniqueId(id);
          doc.setEmbedding(embedding);
          return doc;
      }
  
      private void setPrivateField(Object object, String fieldName, Object value) {
          try {
              java.lang.reflect.Field field = object.getClass().getDeclaredField(fieldName);
              field.setAccessible(true);
              field.set(object, value);
          } catch (Exception e) {
              fail("Failed to set private field: " + fieldName);
          }
      }
  
      private String invokeFilterMethod(String input) {
          try {
              java.lang.reflect.Method method = TinkerPop.class.getDeclaredMethod("filter", String.class);
              method.setAccessible(true);
              return (String) method.invoke(tinkerPop, input);
          } catch (Exception e) {
              fail("Failed to invoke filter method");
              return null;
          }
      }
  }