From e2c4ede405ae21bad4b5b9d361765ede383f9213 Mon Sep 17 00:00:00 2001
From: Elazar Gershuni <elazarg@gmail.com>
Date: Mon, 28 Oct 2024 17:57:41 +0200
Subject: [PATCH] fix tests, add proc test

Signed-off-by: Elazar Gershuni <elazarg@gmail.com>
---
 .gitignore             |  4 ++-
 checkpoint/persist.py  | 59 ++++++++++++++++++++++--------------------
 tests/test_generate.py | 27 ++++++++++++++-----
 3 files changed, 55 insertions(+), 35 deletions(-)

diff --git a/.gitignore b/.gitignore
index 73871f8..ddda751 100644
--- a/.gitignore
+++ b/.gitignore
@@ -7,4 +7,6 @@ __pycache__
 *.log
 *.out
 *.qcow2
-**/cache
\ No newline at end of file
+**/cache
+venv*
+**/dumps
diff --git a/checkpoint/persist.py b/checkpoint/persist.py
index 287852b..7b8e1d4 100644
--- a/checkpoint/persist.py
+++ b/checkpoint/persist.py
@@ -10,7 +10,7 @@
 import hashlib
 import socket
 import struct
-from checkpoint.criu_binding import set_criu, criu_dump
+
 
 FUEL = "FUEL"
 STEP = "STEP"
@@ -243,34 +243,37 @@ def sigint() -> None:
     os.kill(PID, signal.SIGINT)
 
 
-SET_CRIU = False
-CRIU_FOLDER = pathlib.Path("criu_images")
-CRIU_DUMPS = CRIU_FOLDER / "dumps"
+if os.name == "posix":
 
-coredump_iterations = 0
-coredump_steps = 0
+    from checkpoint.criu_binding import set_criu, criu_dump
 
+    SET_CRIU = False
+    CRIU_FOLDER = pathlib.Path("criu_images")
+    CRIU_DUMPS = CRIU_FOLDER / "dumps"
 
-def self_coredump() -> None:
-    global coredump_iterations, coredump_steps
-    if not coredump_iterations:
-        CRIU_FOLDER.mkdir(exist_ok=True)
-        shutil.rmtree(CRIU_DUMPS, ignore_errors=True)
-        CRIU_DUMPS.mkdir(exist_ok=False)
-        set_criu(CRIU_FOLDER)
-    coredump_iterations += 1
+    coredump_iterations = 0
+    coredump_steps = 0
 
-    if coredump_iterations % STEP_VALUE in [0, 1]:
-        criu_dump()
-        image_file = CRIU_FOLDER / "pages-1.img"
-        if not image_file.exists():
-            raise RuntimeError(
-                "CRIU image was not created. Make sure to run the CRIU service:\n"
-                "sudo criu service --shell-job --address /tmp/criu_service.socket"
-            )
-        target_image = CRIU_DUMPS / f"{coredump_steps:05d}.a.img"
-        os.rename(CRIU_FOLDER / "pages-1.img", target_image)
-        if coredump_steps > 0:
-            source_image = CRIU_DUMPS / f"{coredump_steps-1:05d}.b.img"
-            source_image.hardlink_to(target_image)
-        coredump_steps += 1
+    def self_coredump() -> None:
+        global coredump_iterations, coredump_steps
+        if not coredump_iterations:
+            CRIU_FOLDER.mkdir(exist_ok=True)
+            shutil.rmtree(CRIU_DUMPS, ignore_errors=True)
+            CRIU_DUMPS.mkdir(exist_ok=False)
+            set_criu(CRIU_FOLDER)
+        coredump_iterations += 1
+
+        if coredump_iterations % STEP_VALUE in [0, 1]:
+            criu_dump()
+            image_file = CRIU_FOLDER / "pages-1.img"
+            if not image_file.exists():
+                raise RuntimeError(
+                    "CRIU image was not created. Make sure to run the CRIU service:\n"
+                    "sudo criu service --shell-job --address /tmp/criu_service.socket"
+                )
+            target_image = CRIU_DUMPS / f"{coredump_steps:05d}.a.img"
+            os.rename(CRIU_FOLDER / "pages-1.img", target_image)
+            if coredump_steps > 0:
+                source_image = CRIU_DUMPS / f"{coredump_steps-1:05d}.b.img"
+                source_image.hardlink_to(target_image)
+            coredump_steps += 1
diff --git a/tests/test_generate.py b/tests/test_generate.py
index 097f065..711fb85 100644
--- a/tests/test_generate.py
+++ b/tests/test_generate.py
@@ -27,9 +27,7 @@ def naive_transform(filename: pathlib.Path, expected_outfile: pathlib.Path) -> N
     compare_transformed_files(actual, expected_outfile)
 
 
-@pytest.mark.parametrize(
-    "experiment_name", ["k_means", "omp", "pivoter", "trivial"]
-)
+@pytest.mark.parametrize("experiment_name", ["k_means", "omp", "pivoter", "trivial"])
 def test_naive_transformation(experiment_name: str) -> None:
     exp = pathlib.Path("experiment") / experiment_name
     naive_transform(
@@ -45,9 +43,7 @@ def tcp_transform(
     compare_transformed_files(actual, expected_outfile)
 
 
-@pytest.mark.parametrize(
-    "experiment_name", ["k_means", "omp", "pivoter", "trivial"]
-)
+@pytest.mark.parametrize("experiment_name", ["k_means", "omp", "pivoter", "trivial"])
 def test_tcp_transformation(experiment_name: str) -> None:
     exp = pathlib.Path("experiment") / experiment_name
     filename = exp / "main.py"
@@ -60,6 +56,25 @@ def test_tcp_transformation(experiment_name: str) -> None:
     )
 
 
+def proc_transform(
+    filename: pathlib.Path, function_name: str, expected_outfile: pathlib.Path
+) -> None:
+    actual = ast_transform.coredump(filename, function_name)
+    compare_transformed_files(actual, expected_outfile)
+
+
+@pytest.mark.parametrize("experiment_name", ["k_means", "omp", "pivoter", "trivial"])
+def test_proc_transformation(experiment_name: str) -> None:
+    exp = pathlib.Path("experiment") / experiment_name
+    filename = exp / "main.py"
+    expected_outfile = exp / "proc.py"
+    proc_transform(
+        function_name="run",
+        filename=filename,
+        expected_outfile=expected_outfile,
+    )
+
+
 def analyze_and_transform(
     experiment_name: str, function_name: str, simplify: bool
 ) -> None: