From 046931754986e0f53e5c5941abfeb16b34b9201d Mon Sep 17 00:00:00 2001 From: Peng Huo Date: Mon, 5 Aug 2024 17:03:11 -0700 Subject: [PATCH] fix UT Signed-off-by: Peng Huo --- .../opensearch/flint/core/table/OpenSearchIndexTable.scala | 2 +- .../flint/core/table/OpenSearchIndexTableSpec.scala | 4 ++-- 2 files changed, 3 insertions(+), 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 index d016fe9b7..783163687 100644 --- 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 index a30d2de8b..9f6d045fc 100644 --- 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) {