diff --git a/shell.nix b/shell.nix
index fb9d7b7..ba4c573 100644
--- a/shell.nix
+++ b/shell.nix
@@ -6,10 +6,15 @@
 }:
 
 let
-  haskellator = pkgs.callPackage ./default.nix {};
+  haskellator = pkgs.callPackage ./default.nix { };
+  corpus = pkgs.callPackage ./rtlil-corpus/default.nix { };
 in
-pkgs.mkShell {
-  buildInputs = [
-    haskellator
+pkgs.callPackage (
+  { mkShell }:
+  mkShell {
+    CORPUS = corpus;
+    nativeBuildInputs = [
+      haskellator
     ];
-}
+  }
+) { }
diff --git a/test.py b/test.py
index d7f89a4..0dd7505 100644
--- a/test.py
+++ b/test.py
@@ -1,13 +1,30 @@
 import os
 import subprocess
+import argparse
+import shutil
+
+parser = argparse.ArgumentParser(description="Test rtlil-parse")
+parser.add_argument("--corpus", default=os.environ.get('CORPUS'), required=False, help="Path to corpus dir.")
+args = parser.parse_args()
 
 # Define the directories
-corpus_dir = './rtlil-corpus/corpus/'
+if args.corpus:
+    corpus_dir = args.corpus
+else:
+    corpus_dir = './rtlil-corpus/corpus/'
+
 parsed_dir = './parsed_rtlil'
 
 # Create the parsed_rtlil directory if it doesn't exist
 os.makedirs(parsed_dir, exist_ok=True)
 
+# Use rtlil-parse from PATH if available.
+rtil_parse_binary = shutil.which("rtlil-parse")
+if rtil_parse_binary is None:
+    program = "cabal run rtlil-parse --"
+else:
+    program = rtil_parse_binary
+
 # Iterate over all .rtlil files in the corpus directory
 for file_name in os.listdir(corpus_dir):
     if file_name.endswith('.il'):
@@ -16,7 +33,7 @@ for file_name in os.listdir(corpus_dir):
         output_path = os.path.join(parsed_dir, output_file)
 
         # Construct and run the command
-        command = f"cabal run rtlil-parse -- {input_path} {output_path}"
+        command = f"{program} {input_path} {output_path}"
         try:
             subprocess.run(command, shell=True, check=True)
             print(f"Processed: {input_path} -> {output_path}")