Add unit test for segment replication enabled index

Signed-off-by: dreamer-89 <[email protected]>
opensearch-project · Aug 30, 2023 · 8390410 · 8390410
1 parent 9fc9473
commit 8390410
Showing 1 changed file with 33 additions and 0 deletions.
diff --git a/.../java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java b/.../java/org/opensearch/sql/datasources/storage/OpenSearchDataSourceMetadataStorageTest.java
@@ -103,6 +103,39 @@ public void testGetDataSourceMetadata() {
         "basicauth", dataSourceMetadata.getProperties().get("prometheus.auth.type"));
   }
 
+  @SneakyThrows
+  @Test
+  public void testGetDataSourceMetadataWithSegRepEnabled() {
+    Mockito.when(clusterService.state().routingTable().hasIndex(DATASOURCE_INDEX_NAME))
+        .thenReturn(true);
+    Mockito.when(clusterService.state().isSegmentReplicationEnabled(DATASOURCE_INDEX_NAME))
+        .thenReturn(true);
+    Mockito.when(client.search(ArgumentMatchers.any())).thenReturn(searchResponseActionFuture);
+    Mockito.when(searchResponseActionFuture.actionGet()).thenReturn(searchResponse);
+    Mockito.when(searchResponse.status()).thenReturn(RestStatus.OK);
+    Mockito.when(searchResponse.getHits())
+        .thenReturn(
+            new SearchHits(
+                new SearchHit[] {searchHit}, new TotalHits(21, TotalHits.Relation.EQUAL_TO), 1.0F));
+    Mockito.when(searchHit.getSourceAsString()).thenReturn(getBasicDataSourceMetadataString());
+    Mockito.when(encryptor.decrypt("password")).thenReturn("password");
+    Mockito.when(encryptor.decrypt("username")).thenReturn("username");
+
+    Optional<DataSourceMetadata> dataSourceMetadataOptional =
+        openSearchDataSourceMetadataStorage.getDataSourceMetadata(TEST_DATASOURCE_INDEX_NAME);
+
+    Assertions.assertFalse(dataSourceMetadataOptional.isEmpty());
+    DataSourceMetadata dataSourceMetadata = dataSourceMetadataOptional.get();
+    Assertions.assertEquals(TEST_DATASOURCE_INDEX_NAME, dataSourceMetadata.getName());
+    Assertions.assertEquals(DataSourceType.PROMETHEUS, dataSourceMetadata.getConnector());
+    Assertions.assertEquals(
+        "password", dataSourceMetadata.getProperties().get("prometheus.auth.password"));
+    Assertions.assertEquals(
+        "username", dataSourceMetadata.getProperties().get("prometheus.auth.username"));
+    Assertions.assertEquals(
+        "basicauth", dataSourceMetadata.getProperties().get("prometheus.auth.type"));
+  }
+
   @SneakyThrows
   @Test
   public void testGetDataSourceMetadataWith404SearchResponse() {