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 .index .IndexSettings ;
19
+ import org .opensearch .knn .KNNTestCase ;
20
+ import org .opensearch .knn .index .KNNSettings ;
9
21
import org .opensearch .knn .index .VectorDataType ;
10
22
import org .opensearch .knn .index .codec .nativeindex .NativeIndexBuildStrategy ;
11
23
import org .opensearch .knn .index .codec .nativeindex .model .BuildIndexParams ;
16
28
import org .opensearch .knn .index .vectorvalues .TestVectorValues ;
17
29
import org .opensearch .repositories .RepositoriesService ;
18
30
import org .opensearch .repositories .RepositoryMissingException ;
19
- import org .opensearch .test . OpenSearchTestCase ;
31
+ import org .opensearch .repositories . blobstore . BlobStoreRepository ;
20
32
33
+ import java .io .ByteArrayInputStream ;
21
34
import java .io .IOException ;
35
+ import java .io .InputStream ;
22
36
import java .util .List ;
23
37
import java .util .Map ;
38
+ import java .util .Random ;
24
39
25
40
import static org .mockito .ArgumentMatchers .any ;
26
41
import static org .mockito .Mockito .mock ;
27
42
import static org .mockito .Mockito .when ;
43
+ import static org .opensearch .knn .index .KNNSettings .KNN_REMOTE_VECTOR_REPO_SETTING ;
28
44
29
- public class RemoteIndexBuildStrategyTests extends OpenSearchTestCase {
45
+ public class RemoteIndexBuildStrategyTests extends KNNTestCase {
30
46
31
47
static int fallbackCounter = 0 ;
32
48
@@ -38,6 +54,16 @@ public void buildAndWriteIndex(BuildIndexParams indexInfo) throws IOException {
38
54
}
39
55
}
40
56
57
+ @ Before
58
+ @ Override
59
+ public void setUp () throws Exception {
60
+ super .setUp ();
61
+ ClusterSettings clusterSettings = mock (ClusterSettings .class );
62
+ when (clusterSettings .get (KNN_REMOTE_VECTOR_REPO_SETTING )).thenReturn ("test-repo-name" );
63
+ when (clusterService .getClusterSettings ()).thenReturn (clusterSettings );
64
+ KNNSettings .state ().setClusterService (clusterService );
65
+ }
66
+
41
67
public void testFallback () throws IOException {
42
68
List <float []> vectorValues = List .of (new float [] { 1 , 2 }, new float [] { 2 , 3 }, new float [] { 3 , 4 });
43
69
final TestVectorValues .PreDefinedFloatVectorValues randomVectorValues = new TestVectorValues .PreDefinedFloatVectorValues (
@@ -48,7 +74,11 @@ public void testFallback() throws IOException {
48
74
RepositoriesService repositoriesService = mock (RepositoriesService .class );
49
75
when (repositoriesService .repository (any ())).thenThrow (new RepositoryMissingException ("Fallback" ));
50
76
51
- RemoteIndexBuildStrategy objectUnderTest = new RemoteIndexBuildStrategy (() -> repositoriesService , new TestIndexBuildStrategy ());
77
+ RemoteIndexBuildStrategy objectUnderTest = new RemoteIndexBuildStrategy (
78
+ () -> repositoriesService ,
79
+ new TestIndexBuildStrategy (),
80
+ mock (IndexSettings .class )
81
+ );
52
82
53
83
IndexOutputWithBuffer indexOutputWithBuffer = Mockito .mock (IndexOutputWithBuffer .class );
54
84
@@ -64,4 +94,54 @@ public void testFallback() throws IOException {
64
94
objectUnderTest .buildAndWriteIndex (buildIndexParams );
65
95
assertEquals (1 , fallbackCounter );
66
96
}
97
+
98
+ /**
99
+ * Verify the buffered read method in {@link RemoteIndexBuildStrategy#readFromRepository} produces the correct result
100
+ */
101
+ public void testRepositoryRead () throws IOException {
102
+ // Create an InputStream with random values
103
+ int TEST_ARRAY_SIZE = 64 * 1024 * 10 ;
104
+ byte [] byteArray = new byte [TEST_ARRAY_SIZE ];
105
+ Random random = new Random ();
106
+ random .nextBytes (byteArray );
107
+ InputStream randomStream = new ByteArrayInputStream (byteArray );
108
+
109
+ // Create a test segment that we will read/write from
110
+ Directory directory ;
111
+ directory = newFSDirectory (createTempDir ());
112
+ String TEST_SEGMENT_NAME = "test-segment-name" ;
113
+ IndexOutput testIndexOutput = directory .createOutput (TEST_SEGMENT_NAME , IOContext .DEFAULT );
114
+
115
+ // Set up RemoteIndexBuildStrategy and write to IndexOutput
116
+ RepositoriesService repositoriesService = mock (RepositoriesService .class );
117
+ BlobStoreRepository mockRepository = mock (BlobStoreRepository .class );
118
+ BlobPath testBasePath = new BlobPath ().add ("testBasePath" );
119
+ BlobStore mockBlobStore = mock (BlobStore .class );
120
+ AsyncMultiStreamBlobContainer mockBlobContainer = mock (AsyncMultiStreamBlobContainer .class );
121
+
122
+ when (repositoriesService .repository (any ())).thenReturn (mockRepository );
123
+ when (mockRepository .basePath ()).thenReturn (testBasePath );
124
+ when (mockRepository .blobStore ()).thenReturn (mockBlobStore );
125
+ when (mockBlobStore .blobContainer (any ())).thenReturn (mockBlobContainer );
126
+ when (mockBlobContainer .readBlob ("test-blob" )).thenReturn (randomStream );
127
+
128
+ RemoteIndexBuildStrategy objectUnderTest = new RemoteIndexBuildStrategy (
129
+ () -> repositoriesService ,
130
+ mock (NativeIndexBuildStrategy .class ),
131
+ mock (IndexSettings .class )
132
+ );
133
+ // This should read from randomStream into testIndexOutput
134
+ objectUnderTest .readFromRepository ("test-blob" , testIndexOutput );
135
+ testIndexOutput .close ();
136
+
137
+ // Now try to read from the IndexOutput
138
+ IndexInput testIndexInput = directory .openInput (TEST_SEGMENT_NAME , IOContext .DEFAULT );
139
+ byte [] resultByteArray = new byte [TEST_ARRAY_SIZE ];
140
+ testIndexInput .readBytes (resultByteArray , 0 , TEST_ARRAY_SIZE );
141
+ assertArrayEquals (byteArray , resultByteArray );
142
+
143
+ // Test Cleanup
144
+ testIndexInput .close ();
145
+ directory .close ();
146
+ }
67
147
}
0 commit comments