Cleanup filesystem after e2e tests (#1584)

tests/e2e/interactive_mg_runner.py

@@ -80,6 +80,8 @@ ACTIONS = {
     "quit": lambda _: sys.exit(1),
 }
 
+CLEANUP_DIRECTORIES_ON_EXIT = False
+
 log = logging.getLogger("memgraph.tests.e2e")
 
 
@@ -109,10 +111,11 @@ def _start_instance(name, args, log_file, setup_queries, use_ssl, procdir, data_
     assert not is_port_in_use(
         extract_bolt_port(args)
     ), "If this raises, you are trying to start an instance on a port already used by one already running instance."
-    mg_instance = MemgraphInstanceRunner(MEMGRAPH_BINARY, use_ssl)
-    MEMGRAPH_INSTANCES[name] = mg_instance
+
     log_file_path = os.path.join(BUILD_DIR, "logs", log_file)
     data_directory_path = os.path.join(BUILD_DIR, data_directory)
+    mg_instance = MemgraphInstanceRunner(MEMGRAPH_BINARY, use_ssl, {data_directory_path})
+    MEMGRAPH_INSTANCES[name] = mg_instance
     binary_args = args + ["--log-file", log_file_path] + ["--data-directory", data_directory_path]
 
     if len(procdir) != 0:
@@ -122,39 +125,43 @@ def _start_instance(name, args, log_file, setup_queries, use_ssl, procdir, data_
     assert mg_instance.is_running(), "An error occured after starting Memgraph instance: application stopped running."
 
 
-def stop_all():
+def stop_all(keep_directories=True):
     for mg_instance in MEMGRAPH_INSTANCES.values():
-        mg_instance.stop()
+        mg_instance.stop(keep_directories)
     MEMGRAPH_INSTANCES.clear()
 
 
-def stop_instance(context, name):
+def stop_instance(context, name, keep_directories=True):
     for key, _ in context.items():
         if key != name:
             continue
-        MEMGRAPH_INSTANCES[name].stop()
+        MEMGRAPH_INSTANCES[name].stop(keep_directories)
         MEMGRAPH_INSTANCES.pop(name)
 
 
-def stop(context, name):
+def stop(context, name, keep_directories=True):
     if name != "all":
-        stop_instance(context, name)
+        stop_instance(context, name, keep_directories)
         return
 
     stop_all()
 
 
-def kill(context, name):
+def kill(context, name, keep_directories=True):
     for key in context.keys():
         if key != name:
             continue
-        MEMGRAPH_INSTANCES[name].kill()
+        MEMGRAPH_INSTANCES[name].kill(keep_directories)
         MEMGRAPH_INSTANCES.pop(name)
 
 
+def cleanup_directories_on_exit(value=True):
+    CLEANUP_DIRECTORIES_ON_EXIT = value
+
+
 @atexit.register
 def cleanup():
-    stop_all()
+    stop_all(CLEANUP_DIRECTORIES_ON_EXIT)
 
 
 def start_instance(context, name, procdir):
@@ -184,8 +191,8 @@ def start_instance(context, name, procdir):
     assert len(mg_instances) == 1
 
 
-def start_all(context, procdir=""):
-    stop_all()
+def start_all(context, procdir="", keep_directories=True):
+    stop_all(keep_directories)
     for key, _ in context.items():
         start_instance(context, key, procdir)
 

tests/e2e/memgraph.py

@@ -11,6 +11,7 @@
 
 import copy
 import os
+import shutil
 import subprocess
 import sys
 import time
@@ -56,13 +57,14 @@ def replace_paths(path):
 
 
 class MemgraphInstanceRunner:
-    def __init__(self, binary_path=MEMGRAPH_BINARY, use_ssl=False):
+    def __init__(self, binary_path=MEMGRAPH_BINARY, use_ssl=False, delete_on_stop=None):
         self.host = "127.0.0.1"
         self.bolt_port = None
         self.binary_path = binary_path
         self.args = None
         self.proc_mg = None
         self.ssl = use_ssl
+        self.delete_on_stop = delete_on_stop
 
     def execute_setup_queries(self, setup_queries):
         if setup_queries is None:
@@ -128,7 +130,7 @@ class MemgraphInstanceRunner:
             return False
         return True
 
-    def stop(self):
+    def stop(self, keep_directories=False):
         if not self.is_running():
             return
 
@@ -140,9 +142,16 @@ class MemgraphInstanceRunner:
 
         time.sleep(1)
 
-    def kill(self):
+        if not keep_directories:
+            for folder in self.delete_on_stop or {}:
+                shutil.rmtree(folder)
+
+    def kill(self, keep_directories=False):
         if not self.is_running():
             return
         self.proc_mg.kill()
         code = self.proc_mg.wait()
+        if not keep_directories:
+            for folder in self.delete_on_stop or {}:
+                shutil.rmtree(folder)
         assert code == -9, "The killed Memgraph process exited with non-nine!"

tests/e2e/runner.py

@@ -53,8 +53,8 @@ def run(args):
 
         # Setup.
         @atexit.register
-        def cleanup():
-            interactive_mg_runner.stop_all()
+        def cleanup(keep_directories=True):
+            interactive_mg_runner.stop_all(keep_directories)
 
         if "pre_set_workload" in workload:
             binary = os.path.join(BUILD_DIR, workload["pre_set_workload"])
@@ -92,7 +92,7 @@ def run(args):
                     data = mg_instance.query(validation["query"], conn)[0][0]
                     assert data == validation["expected"]
                 conn.close()
-        cleanup()
+        cleanup(keep_directories=False)
         log.info("%s PASSED.", workload_name)