From 406ed7eb4d15160c5ab7e4f28d3c790810150f14 Mon Sep 17 00:00:00 2001
From: Aneesh Durg <aneeshdurg17@gmail.com>
Date: Sat, 7 Sep 2024 15:36:10 -0500
Subject: [PATCH] Add test

---
 .github/workflows/test.yml | 36 ++++++++++++++++++++++++++++++++++++
 test/replace_return.py     | 17 +++++++++++++++++
 test/test_module_pass.py   | 30 ++++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+)
 create mode 100644 .github/workflows/test.yml
 create mode 100644 test/replace_return.py
 create mode 100644 test/test_module_pass.py

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
new file mode 100644
index 0000000..02e7374
--- /dev/null
+++ b/.github/workflows/test.yml
@@ -0,0 +1,36 @@
+name: Test
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  test:
+    name: Test PyLLVMPass
+    runs-on: ubuntu-latest
+    steps:
+      - name: "checkout repository"
+        uses: actions/checkout@v4
+
+      - name: "install rust"
+        uses: actions-rust-lang/setup-rust-toolchain@v1
+
+      - name: "install llvm-18"
+        run: |
+          wget https://apt.llvm.org/llvm.sh
+          chmod u+x llvm.sh
+          sudo ./llvm.sh 18
+
+      - name: "install pytest"
+        run: |
+          pip install pytest
+
+      - name: "build"
+        run: |
+          LLVM_CONFIG=llvm-config-18 cargo build
+
+      - name: "run tests"
+        run: |
+          cd test
+          LLVM_CONFIG=llvm-config-18 PYLLVM_PLUGIN_PATH=../target/debug/libpyllvmpass.so pytest
diff --git a/test/replace_return.py b/test/replace_return.py
new file mode 100644
index 0000000..95bc0b8
--- /dev/null
+++ b/test/replace_return.py
@@ -0,0 +1,17 @@
+"""LLVM Module pass that replace all return statements with `return 1`"""
+
+import llvmcpy.llvm as cllvm
+
+def run_on_module(module: cllvm.Module):
+    ctx = module.get_context()
+    i32_t = ctx.int32_type()
+    # constant i32(1)
+    const_1 = i32_t.const_int(1, False)
+
+    for fn in module.iter_functions():
+        for bb in fn.iter_basic_blocks():
+            for inst in bb.iter_instructions():
+                if inst.instruction_opcode != cllvm.Ret:
+                    continue
+                inst.set_arg_operand(0, const_1)
+    return 1
diff --git a/test/test_module_pass.py b/test/test_module_pass.py
new file mode 100644
index 0000000..96e8def
--- /dev/null
+++ b/test/test_module_pass.py
@@ -0,0 +1,30 @@
+import os
+import subprocess
+import pytest
+
+@pytest.fixture
+def plugin_path():
+    return os.environ["PYLLVM_PLUGIN_PATH"]
+
+
+def test_replace_return(tmp_path, plugin_path):
+    src = tmp_path / 'main.cpp'
+    with src.open('w') as f:
+        f.write("int main() { return 0; }\n")
+    ll = tmp_path/'main.ll'
+    subprocess.check_call(['clang-18', '-O3', src, '-S', '-emit-llvm', '-o', ll])
+    assert 'ret i32 0' in ll.read_text()
+
+    outfile = tmp_path/'out.ll'
+    subprocess.check_call([
+        'opt-18',
+        '--load-pass-plugin',
+        plugin_path,
+        '--passes=pyllvmpass[replace_return]',
+        ll,
+        '-S',
+        '-o',
+        outfile
+    ])
+    assert 'ret i32 0' not in outfile.read_text()
+    assert 'ret i32 1' in outfile.read_text()