embedding space component tests

theam · Aug 31, 2023 · 89c8bf7 · 89c8bf7
1 parent cb5aabd
commit 89c8bf7
Showing 1 changed file with 40 additions and 8 deletions.
diff --git a/...test/java/com/theagilemonkeys/ellmental/embeddingsspace/EmbeddingsSpaceComponentTest.java b/...test/java/com/theagilemonkeys/ellmental/embeddingsspace/EmbeddingsSpaceComponentTest.java
@@ -6,14 +6,12 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 import org.mockito.*;
-import org.mockito.junit.MockitoJUnitRunner;
 import org.mockito.junit.jupiter.MockitoExtension;
 
 import java.util.List;
 import java.util.UUID;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
-import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.*;
 
 @ExtendWith(MockitoExtension.class)
@@ -29,12 +27,13 @@ public class EmbeddingsSpaceComponentTest {
     private ArgumentCaptor<String> sampleTextCaptor;
 
     private final String sampleText = "sample text";
+    private final UUID embeddingId = UUID.randomUUID();
+    private final Embedding embeddingMock = new Embedding(embeddingId, List.of(1.0, 2.0, 3.0), null);
 
     @Test
     public void testSaveEmbedding() {
-        Embedding embeddingMock = new Embedding(UUID.randomUUID(), List.of(1.0, 2.0, 3.0), null);
         when(embeddingsGenerationModel.generateEmbedding(sampleText)).thenReturn(embeddingMock);
-        //doNothing().when(embeddingsStore).store(embeddingMock);
+
         embeddingsSpaceComponent.save(sampleText);
 
         verify(embeddingsGenerationModel).generateEmbedding(sampleTextCaptor.capture());
@@ -48,21 +47,54 @@ public void testSaveEmbedding() {
 
     @Test
     public void testMostSimilarEmbeddings() {
-        assertTrue(true);
+        int limit = 1;
+        Embedding similarEmbeddingMock = new Embedding(UUID.randomUUID(), List.of(4.0, 5.0, 6.0), null);
+
+        when(embeddingsStore.similaritySearch(embeddingMock, limit)).thenReturn(List.of(similarEmbeddingMock));
+
+        List<Embedding> embeddings = embeddingsSpaceComponent.mostSimilarEmbeddings(embeddingMock, limit);
+
+        assertEquals(embeddings.get(0), similarEmbeddingMock);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
     }
 
     @Test
     public void testMostSimilarEmbeddingsWhenUsingString() {
-        assertTrue(true);
+        int limit = 1;
+        Embedding similarEmbeddingMock = new Embedding(UUID.randomUUID(), List.of(4.0, 5.0, 6.0), null);
+
+        when(embeddingsGenerationModel.generateEmbedding(sampleText)).thenReturn(embeddingMock);
+        when(embeddingsStore.similaritySearch(embeddingMock, limit)).thenReturn(List.of(similarEmbeddingMock));
+
+        List<Embedding> embeddings = embeddingsSpaceComponent.mostSimilarEmbeddings(sampleText, limit);
+
+        assertEquals(embeddings.get(0), similarEmbeddingMock);
+
+        verify(embeddingsGenerationModel).generateEmbedding(sampleTextCaptor.capture());
+        String capturedSampleText = sampleTextCaptor.getValue();
+        assertEquals(sampleText, capturedSampleText);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
     }
 
     @Test
     public void testGetEmbedding() {
-        assertTrue(true);
+        when(embeddingsStore.get(embeddingId)).thenReturn(embeddingMock);
+
+        Embedding embedding = embeddingsSpaceComponent.get(embeddingId);
+
+        assertEquals(embeddingMock, embedding);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
     }
 
     @Test
     public void testDeleteEmbedding() {
-        assertTrue(true);
+        embeddingsSpaceComponent.delete(embeddingId);
+
+        verify(embeddingsStore).delete(embeddingId);
+
+        verifyNoMoreInteractions(embeddingsGenerationModel, embeddingsStore);
     }
 }