test(GDS): fix num_batches issue

tigergraph · Nov 10, 2023 · f7dd719 · f7dd719
1 parent 2e565a4
commit f7dd719
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/tests/test_gds_GDS.py b/tests/test_gds_GDS.py
@@ -26,7 +26,7 @@ def test_neighborLoader(self):
             buffer_size=4,
         )
         self.assertTrue(is_query_installed(self.conn, loader.query_name))
-        self.assertEqual(loader.num_batches, 9)
+        self.assertEqual(loader.batch_size, 16)
 
     def test_neighborLoader_multiple_filters(self):
         loaders = self.conn.gds.neighborLoader(
@@ -60,7 +60,7 @@ def test_graphLoader(self):
             buffer_size=4,
         )
         self.assertTrue(is_query_installed(self.conn, loader.query_name))
-        self.assertEqual(loader.num_batches, 11)
+        self.assertEqual(loader.batch_size, 1024)
 
     def test_vertexLoader(self):
         loader = self.conn.gds.vertexLoader(
@@ -72,7 +72,7 @@ def test_vertexLoader(self):
             buffer_size=4,
         )
         self.assertTrue(is_query_installed(self.conn, loader.query_name))
-        self.assertEqual(loader.num_batches, 9)
+        self.assertEqual(loader.batch_size, 16)
 
     def test_edgeLoader(self):
         loader = self.conn.gds.edgeLoader(
@@ -83,7 +83,7 @@ def test_edgeLoader(self):
             buffer_size=4,
         )
         self.assertTrue(is_query_installed(self.conn, loader.query_name))
-        self.assertEqual(loader.num_batches, 11)
+        self.assertEqual(loader.batch_size, 1024)
 
     def test_edgeNeighborLoader(self):
         loader = self.conn.gds.edgeNeighborLoader(
@@ -100,7 +100,7 @@ def test_edgeNeighborLoader(self):
             buffer_size=4,
         )
         self.assertTrue(is_query_installed(self.conn, loader.query_name))
-        self.assertEqual(loader.num_batches, 11)
+        self.assertEqual(loader.batch_size, 1024)
 
     def test_configureKafka(self):
         self.conn.gds.configureKafka(kafka_address="kafka:9092")