Refactor mgbench

tests/mgbench/benchmark_context.py

@@ -9,8 +9,13 @@
 # by the Apache License, Version 2.0, included in the file
 # licenses/APL.txt.
 
+from workload_mode import (
+    BENCHMARK_MODE_ISOLATED,
+    BENCHMARK_MODE_MIXED,
+    BENCHMARK_MODE_REALISTIC,
+)
+
 
-# Describes all the information of single benchmark.py run.
 class BenchmarkContext:
     """
     Class for holding information on what type of benchmark is being executed
@@ -38,6 +43,8 @@ class BenchmarkContext:
         no_authorization: bool = True,
         customer_workloads: str = None,
         vendor_args: dict = {},
+        disk_storage: bool = False,
+        in_memory_analytical: bool = False,
     ) -> None:
         self.benchmark_target_workload = benchmark_target_workload
         self.vendor_binary = vendor_binary
@@ -52,15 +59,19 @@ class BenchmarkContext:
         self.export_results = export_results
         self.temporary_directory = temporary_directory
 
+        assert (
+            workload_mixed is None or workload_realistic is None
+        ), "Cannot run both mixed and realistic workload, please select one!"
+
         if workload_mixed != None:
-            self.mode = "Mixed"
+            self.mode = BENCHMARK_MODE_MIXED
             self.mode_config = workload_mixed
         elif workload_realistic != None:
-            self.mode = "Realistic"
+            self.mode = BENCHMARK_MODE_REALISTIC
             self.mode_config = workload_realistic
         else:
-            self.mode = "Isolated"
-            self.mode_config = "Isolated run does not have a config."
+            self.mode = BENCHMARK_MODE_ISOLATED
+            self.mode_config = None
 
         self.time_dependent_execution = time_dependent_execution
         self.performance_tracking = performance_tracking
@@ -70,6 +81,8 @@ class BenchmarkContext:
         self.vendor_args = vendor_args
         self.active_workload = None
         self.active_variant = None
+        self.disk_storage = disk_storage
+        self.in_memory_analytical = in_memory_analytical
 
     def set_active_workload(self, workload: str) -> None:
         self.active_workload = workload

tests/mgbench/graph_bench.py

@@ -127,6 +127,8 @@ def run_full_benchmarks(
         ],
     ]
 
+    assert not realistic or not mixed, "Cannot run both realistic and mixed workload, please select one!"
+
     if realistic:
         # Configurations for full workload
         for count, write, read, update, analytical in realistic:

tests/mgbench/helpers.py

@@ -188,6 +188,7 @@ def filter_workloads(available_workloads: dict, benchmark_context: BenchmarkCont
             raise Exception("Invalid benchmark description '" + pattern + "'!")
         pattern.extend(["", "*", "*"][len(pattern) - 1 :])
         patterns[i] = pattern
+
     filtered = []
     for workload in sorted(available_workloads.keys()):
         generator, queries = available_workloads[workload]

tests/mgbench/workload_mode.py

@@ -0,0 +1,14 @@
+# Copyright 2023 Memgraph Ltd.
+#
+# Use of this software is governed by the Business Source License
+# included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
+# License, and you may not use this file except in compliance with the Business Source License.
+#
+# As of the Change Date specified in that file, in accordance with
+# the Business Source License, use of this software will be governed
+# by the Apache License, Version 2.0, included in the file
+# licenses/APL.txt.
+
+BENCHMARK_MODE_MIXED = "Mixed"
+BENCHMARK_MODE_REALISTIC = "Realistic"
+BENCHMARK_MODE_ISOLATED = "Isolated"