From baacc52a65fbb373bca79a390a7334abb0580113 Mon Sep 17 00:00:00 2001
From: jeremy <jeremy.bailleux@memgraph.io>
Date: Mon, 7 Nov 2022 09:54:14 +0100
Subject: [PATCH] Add support for split file configuration

---
 tests/mgbench/benchmark.py |  6 +++++-
 tests/mgbench/datasets.py  | 15 +++++++++++++++
 tests/mgbench/runners.py   |  5 ++++-
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/tests/mgbench/benchmark.py b/tests/mgbench/benchmark.py
index 44eb2290d..40760f63e 100755
--- a/tests/mgbench/benchmark.py
+++ b/tests/mgbench/benchmark.py
@@ -181,7 +181,11 @@ for dataset, tests in benchmarks:
 
     # Prepare runners and import the dataset.
     memgraph = runners.Memgraph(
-        args.memgraph_binary, args.temporary_directory, not args.no_properties_on_edges, args.test_system_args
+        args.memgraph_binary,
+        args.temporary_directory,
+        not args.no_properties_on_edges,
+        args.test_system_args,
+        dataset.get_split_file(),
     )
     client = runners.Client(args.client_binary, args.temporary_directory)
     memgraph.start_preparation()
diff --git a/tests/mgbench/datasets.py b/tests/mgbench/datasets.py
index 3a5806629..e73169aab 100644
--- a/tests/mgbench/datasets.py
+++ b/tests/mgbench/datasets.py
@@ -63,6 +63,10 @@ class Dataset:
             raise ValueError("The size defined for this variant doesn't " "have the number of vertices and/or edges!")
         self._num_vertices = self._size["vertices"]
         self._num_edges = self._size["edges"]
+        if self.SPLIT_FILES is not None:
+            self._split_file = self.SPLIT_FILES.get(variant, None)
+        else:
+            self._split_file = None
 
     def prepare(self, directory):
         if self._file is not None:
@@ -92,6 +96,11 @@ class Dataset:
         """Returns number of vertices/edges for the current variant."""
         return self._size
 
+    def get_split_file(self):
+        """Returns the location of the split file of the dataset."""
+        assert self._split_file is not None
+        return self._split_file
+
     # All tests should be query generator functions that output all of the
     # queries that should be executed by the runner. The functions should be
     # named `benchmark__GROUPNAME__TESTNAME` and should not accept any
@@ -299,6 +308,12 @@ class AccessControl(Dataset):
         "medium": "https://s3.eu-west-1.amazonaws.com/deps.memgraph.io/dataset/accesscontrol/accesscontrol_medium.setup.cypher.gz",
         "large": "https://s3.eu-west-1.amazonaws.com/deps.memgraph.io/dataset/accesscontrol/accesscontrol_large.setup.cypher.gz",
     }
+    SPLIT_FILES = {
+        "empty_only_index": "splitfiles/accesscontrol_small.shard_configuration",
+        "small": "splitfiles/accesscontrol_small.shard_configuration",
+        "medium": "splitfiles/accesscontrol_medium.shard_configuration",
+        "large": "splitfiles/accesscontrol_large.shard_configuration",
+    }
     SIZES = {
         "empty_only_index": {
             "vertices": 0,
diff --git a/tests/mgbench/runners.py b/tests/mgbench/runners.py
index 2b69a811f..acee68e07 100644
--- a/tests/mgbench/runners.py
+++ b/tests/mgbench/runners.py
@@ -51,12 +51,13 @@ def _get_usage(pid):
 
 
 class Memgraph:
-    def __init__(self, memgraph_binary, temporary_dir, properties_on_edges, extra_args):
+    def __init__(self, memgraph_binary, temporary_dir, properties_on_edges, extra_args, split_file):
         self._memgraph_binary = memgraph_binary
         self._directory = tempfile.TemporaryDirectory(dir=temporary_dir)
         self._properties_on_edges = properties_on_edges
         self._proc_mg = None
         self._extra_args = extra_args
+        self._split_file = split_file
         atexit.register(self._cleanup)
 
         # Determine Memgraph version
@@ -85,6 +86,8 @@ class Memgraph:
             for i in range(0, len(args_list), 2):
                 kwargs[args_list[i]] = args_list[i + 1]
 
+        kwargs["split-file"] = self._split_file
+
         return _convert_args_to_flags(self._memgraph_binary, **kwargs)
 
     def _start(self, **kwargs):