5
5
6
6
package org .opensearch .knn .index .codec .nativeindex .remote ;
7
7
8
+ import org .apache .lucene .store .Directory ;
9
+ import org .apache .lucene .store .IOContext ;
10
+ import org .apache .lucene .store .IndexInput ;
11
+ import org .apache .lucene .store .IndexOutput ;
12
+ import org .junit .Before ;
8
13
import org .mockito .Mockito ;
14
+ import org .opensearch .common .blobstore .AsyncMultiStreamBlobContainer ;
15
+ import org .opensearch .common .blobstore .BlobPath ;
16
+ import org .opensearch .common .blobstore .BlobStore ;
17
+ import org .opensearch .common .settings .ClusterSettings ;
18
+ import org .opensearch .knn .KNNTestCase ;
19
+ import org .opensearch .knn .index .KNNSettings ;
9
20
import org .opensearch .knn .index .VectorDataType ;
10
21
import org .opensearch .knn .index .codec .nativeindex .NativeIndexBuildStrategy ;
11
22
import org .opensearch .knn .index .codec .nativeindex .model .BuildIndexParams ;
16
27
import org .opensearch .knn .index .vectorvalues .TestVectorValues ;
17
28
import org .opensearch .repositories .RepositoriesService ;
18
29
import org .opensearch .repositories .RepositoryMissingException ;
19
- import org .opensearch .test . OpenSearchTestCase ;
30
+ import org .opensearch .repositories . blobstore . BlobStoreRepository ;
20
31
32
+ import java .io .ByteArrayInputStream ;
21
33
import java .io .IOException ;
34
+ import java .io .InputStream ;
22
35
import java .util .List ;
23
36
import java .util .Map ;
37
+ import java .util .Random ;
24
38
25
39
import static org .mockito .ArgumentMatchers .any ;
26
40
import static org .mockito .Mockito .mock ;
27
41
import static org .mockito .Mockito .when ;
42
+ import static org .opensearch .knn .index .KNNSettings .KNN_REMOTE_VECTOR_REPO_SETTING ;
28
43
29
- public class RemoteIndexBuildStrategyTests extends OpenSearchTestCase {
44
+ public class RemoteIndexBuildStrategyTests extends KNNTestCase {
30
45
31
46
static int fallbackCounter = 0 ;
32
47
@@ -38,6 +53,16 @@ public void buildAndWriteIndex(BuildIndexParams indexInfo) throws IOException {
38
53
}
39
54
}
40
55
56
+ @ Before
57
+ @ Override
58
+ public void setUp () throws Exception {
59
+ super .setUp ();
60
+ ClusterSettings clusterSettings = mock (ClusterSettings .class );
61
+ when (clusterSettings .get (KNN_REMOTE_VECTOR_REPO_SETTING )).thenReturn ("test-repo-name" );
62
+ when (clusterService .getClusterSettings ()).thenReturn (clusterSettings );
63
+ KNNSettings .state ().setClusterService (clusterService );
64
+ }
65
+
41
66
public void testFallback () throws IOException {
42
67
List <float []> vectorValues = List .of (new float [] { 1 , 2 }, new float [] { 2 , 3 }, new float [] { 3 , 4 });
43
68
final TestVectorValues .PreDefinedFloatVectorValues randomVectorValues = new TestVectorValues .PreDefinedFloatVectorValues (
@@ -64,4 +89,68 @@ public void testFallback() throws IOException {
64
89
objectUnderTest .buildAndWriteIndex (buildIndexParams );
65
90
assertEquals (1 , fallbackCounter );
66
91
}
92
+
93
+ /**
94
+ * Verify the buffered read method in {@link RemoteIndexBuildStrategy#readFromRepository} produces the correct result
95
+ */
96
+ public void testRepositoryRead () throws IOException {
97
+ String TEST_FILE_NAME = randomAlphaOfLength (8 ) + KNNEngine .FAISS .getExtension ();
98
+
99
+ // Create an InputStream with random values
100
+ int TEST_ARRAY_SIZE = 64 * 1024 * 10 ;
101
+ byte [] byteArray = new byte [TEST_ARRAY_SIZE ];
102
+ Random random = new Random ();
103
+ random .nextBytes (byteArray );
104
+ InputStream randomStream = new ByteArrayInputStream (byteArray );
105
+
106
+ // Create a test segment that we will read/write from
107
+ Directory directory ;
108
+ directory = newFSDirectory (createTempDir ());
109
+ String TEST_SEGMENT_NAME = "test-segment-name" ;
110
+ IndexOutput testIndexOutput = directory .createOutput (TEST_SEGMENT_NAME , IOContext .DEFAULT );
111
+ IndexOutputWithBuffer testIndexOutputWithBuffer = new IndexOutputWithBuffer (testIndexOutput );
112
+
113
+ // Set up RemoteIndexBuildStrategy and write to IndexOutput
114
+ RepositoriesService repositoriesService = mock (RepositoriesService .class );
115
+ BlobStoreRepository mockRepository = mock (BlobStoreRepository .class );
116
+ BlobPath testBasePath = new BlobPath ().add ("testBasePath" );
117
+ BlobStore mockBlobStore = mock (BlobStore .class );
118
+ AsyncMultiStreamBlobContainer mockBlobContainer = mock (AsyncMultiStreamBlobContainer .class );
119
+
120
+ when (repositoriesService .repository (any ())).thenReturn (mockRepository );
121
+ when (mockRepository .basePath ()).thenReturn (testBasePath );
122
+ when (mockRepository .blobStore ()).thenReturn (mockBlobStore );
123
+ when (mockBlobStore .blobContainer (any ())).thenReturn (mockBlobContainer );
124
+ when (mockBlobContainer .readBlob (TEST_FILE_NAME )).thenReturn (randomStream );
125
+
126
+ RemoteIndexBuildStrategy objectUnderTest = new RemoteIndexBuildStrategy (
127
+ () -> repositoriesService ,
128
+ mock (NativeIndexBuildStrategy .class )
129
+ );
130
+
131
+ // Verify file extension check
132
+ assertThrows (IllegalArgumentException .class , () -> objectUnderTest .readFromRepository ("test_file.txt" , testIndexOutputWithBuffer ));
133
+
134
+ // Now test with valid file extensions
135
+ String testPath = randomFrom (
136
+ List .of (
137
+ "testBasePath/testDirectory/" + TEST_FILE_NAME , // Test with subdirectory
138
+ "testBasePath/" + TEST_FILE_NAME , // Test with only base path
139
+ TEST_FILE_NAME // test with no base path
140
+ )
141
+ );
142
+ // This should read from randomStream into testIndexOutput
143
+ objectUnderTest .readFromRepository (testPath , testIndexOutputWithBuffer );
144
+ testIndexOutput .close ();
145
+
146
+ // Now try to read from the IndexOutput
147
+ IndexInput testIndexInput = directory .openInput (TEST_SEGMENT_NAME , IOContext .DEFAULT );
148
+ byte [] resultByteArray = new byte [TEST_ARRAY_SIZE ];
149
+ testIndexInput .readBytes (resultByteArray , 0 , TEST_ARRAY_SIZE );
150
+ assertArrayEquals (byteArray , resultByteArray );
151
+
152
+ // Test Cleanup
153
+ testIndexInput .close ();
154
+ directory .close ();
155
+ }
67
156
}
0 commit comments