fix UT

Signed-off-by: Peng Huo <[email protected]>
opensearch-project · Aug 6, 2024 · 0469317 · 0469317
1 parent d270b41
commit 0469317
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/flint-core/src/main/scala/org/opensearch/flint/core/table/OpenSearchIndexTable.scala b/flint-core/src/main/scala/org/opensearch/flint/core/table/OpenSearchIndexTable.scala
@@ -57,7 +57,7 @@ class OpenSearchIndexTable(metaData: MetaData, option: FlintOptions) extends Tab
       if (docCount == 0) {
         maxResultWindow
       } else {
-        val totalSizeBytes = indexStats.total().store().sizeInBytes
+        val totalSizeBytes = indexStats.primaries().store().sizeInBytes
         val docSize = Math.ceil(totalSizeBytes / docCount).toLong
         Math.max(Math.min(maxSplitSizeBytes / docSize, maxResultWindow), 1).toInt
       }

diff --git a/flint-core/src/test/scala/org/opensearch/flint/core/table/OpenSearchIndexTableSpec.scala b/flint-core/src/test/scala/org/opensearch/flint/core/table/OpenSearchIndexTableSpec.scala
@@ -53,9 +53,9 @@ class OpenSearchIndexTableSpec
     when(options.supportShard()).thenReturn(supportShard)
 
     when(openSearchClient.stats(any[IndicesStatsRequest])).thenReturn(mockIndicesStatsResp)
-    when(mockIndicesStatsResp.indices().get(any[String]).total().docs().count())
+    when(mockIndicesStatsResp.indices().get(any[String]).primaries().docs().count())
       .thenReturn(docCount)
-    when(mockIndicesStatsResp.indices().get(any[String]).total().store().sizeInBytes)
+    when(mockIndicesStatsResp.indices().get(any[String]).primaries().store().sizeInBytes)
       .thenReturn(storeSizeInBytes)
 
     new OpenSearchIndexTable(metaData, options) {