ATLAS-5147: Relationship Search API approximateCount returns pageLimit value instead of total relationship count#2

Author: aditya-gupta36

graphdb/janus/src/main/java/org/apache/atlas/repository/graphdb/janus/AtlasJanusIndexQuery.java

@@ -114,7 +114,7 @@ public Iterator<Result<AtlasJanusVertex, AtlasJanusEdge>> edges(int offset, int
 
     @Override
     public Long edgeTotals() {
-        return query.edgeTotals();
+        return query.offset(0).limit(Integer.MAX_VALUE).edgeTotals();
     }
 
     /**

repository/src/test/java/org/apache/atlas/discovery/EntitySearchProcessorTest.java

@@ -40,6 +40,7 @@
 import org.slf4j.LoggerFactory;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
@@ -836,4 +837,21 @@ private void createChineseEntityWithDescription() throws AtlasBaseException {
 
         cjkGuid2 = entityHeader.getGuid();
     }
+
+    @DataProvider(name = "limitProvider")
+    public Integer[] limitProvider() {
+        return new Integer[] {15, 0, 1, 10, 20};
+    }
+
+    @Test(dependsOnMethods = "searchWithNEQ_stringAttr", dataProvider = "limitProvider")
+    public void testEntitySearchApproxCountForAllLimits(int limit) throws AtlasBaseException {
+        SearchParameters params = new SearchParameters();
+        params.setTypeName(HIVE_TABLE_TYPE);
+        params.setLimit(limit);
+
+        SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.emptySet());
+        EntitySearchProcessor processor = new EntitySearchProcessor(context);
+
+        assertEquals(processor.getResultCount(), 11);
+    }
 }

repository/src/test/java/org/apache/atlas/discovery/RelationshipSearchProcessorTest.java

@@ -29,6 +29,7 @@
 import org.apache.commons.collections.CollectionUtils;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
+import org.testng.annotations.DataProvider;
 import org.testng.annotations.Guice;
 import org.testng.annotations.Test;
 
@@ -65,6 +66,23 @@ public void totalRelationships() throws AtlasBaseException {
         executeAndAssert(params, 17);
     }
 
+    @DataProvider(name = "limitProvider")
+    public Integer[] limitProvider() {
+        return new Integer[] {15, 0, 1, 17, 20};
+    }
+
+    @Test(dataProvider = "limitProvider")
+    public void testRelationshipsApproxCountForAllLimits(int limit) throws AtlasBaseException {
+        SearchParameters params = new SearchParameters();
+        params.setRelationshipName(USER_POST_TYPE);
+        params.setLimit(limit);
+
+        SearchContext context = new SearchContext(params, typeRegistry, graph, Collections.emptySet());
+        RelationshipSearchProcessor processor = new RelationshipSearchProcessor(context, indexer.getEdgeIndexKeys());
+
+        assertEquals(processor.getResultCount(), 17);
+    }
+
     @Test
     public void searchByAttribute() throws AtlasBaseException {
         SearchParameters params = new SearchParameters();