diff --git a/src/test.cpp b/src/test.cpp
index db35882..75a8ef4 100644
--- a/src/test.cpp
+++ b/src/test.cpp
@@ -1,6 +1,11 @@
 #include "QuantumComputerSimulator.h"
 #include <iostream>
 
+#include <string>   // Needed for std::string
+#include <vector>   // Needed for std::vector (though maybe already included)
+#include <stdexcept> // Needed for std::stoi exception handling (optional but good)
+
+
 using namespace std;
 
 void _test_random() {
@@ -251,16 +256,88 @@ void _test_Shor_factorization() {
 }
 
 
-int main() {
 
-	_test_random();
-	_test_unitary();
-	_test_collapse();
-	_test_QFT();
-	_test_quantum_add();
-	_test_Grover();
-	_test_period_find();
-	_test_Shor_factorization();
-    
-	return 0;
+int main(int argc, char *argv[]) {
+
+    // Default behavior: If no arguments are provided, run all tests
+    if (argc == 1) {
+        cout << "No arguments provided. Running all tests..." << endl;
+        _test_random();
+        _test_unitary();
+        _test_collapse();
+        _test_QFT();
+        _test_quantum_add();
+        _test_Grover();
+        _test_period_find();
+        _test_Shor_factorization();
+        cout << "\nFinished running all tests." << endl;
+        return 0;
+    }
+
+    // --- Argument Parsing ---
+    string command = argv[1]; // The first argument is the command
+
+    try { // Use try-catch for stoul robustness
+        if (command == "random") {
+            _test_random();
+        } else if (command == "unitary") {
+            _test_unitary();
+        } else if (command == "collapse") {
+            _test_collapse();
+        } else if (command == "qft") {
+            _test_QFT();
+        } else if (command == "add") {
+            _test_quantum_add();
+        } else if (command == "grover") {
+            if (argc != 4) {
+                cerr << "Usage: ./test grover <omega> <num_bits>" << endl;
+                return 1;
+            }
+            unsigned int omega = stoul(argv[2]); // Use stoul for unsigned int
+            unsigned int num_bits = stoul(argv[3]);
+            cout << "Testing Grover: Searching for state " << omega << " using " << num_bits << " qubits..." << endl;
+            unsigned int result = Grover(omega, num_bits);
+             cout << "\tGrover result: want " << omega << ", got " << result << endl;
+            cout << "Finished Grover test." << endl;
+        } else if (command == "period_find") {
+            if (argc != 4) {
+                cerr << "Usage: ./test period_find <a> <N>" << endl;
+                return 1;
+            }
+            unsigned int a = stoul(argv[2]);
+            unsigned int N = stoul(argv[3]);
+            cout << "Testing Period Finding: Finding period of " << a << "^x mod " << N << "..." << endl;
+            unsigned int r = find_Shor_period(a, N);
+             cout << "\tQuantum period find found the period of " << a << "^x mod " << N << " to be " << r << endl;
+             cout << "\tCheck: " << a << "^" << r << " = " << mod_power(a, r, N) << " mod " << N << endl;
+            cout << "Finished period finding test." << endl;
+        } else if (command == "shor") {
+             if (argc != 3) {
+                cerr << "Usage: ./test shor <N>" << endl;
+                return 1;
+            }
+            unsigned int N = stoul(argv[2]);
+             cout << "Testing Shor: Factorizing N = " << N << "..." << endl;
+            unsigned int factor = Shor(N);
+            cout << "\tShor algorithm found a factor of " << N << " to be " << factor << endl;
+            cout << "Finished Shor factorization test." << endl;
+        } else {
+            cerr << "Unknown command: " << command << endl;
+            cerr << "Available commands: random, unitary, collapse, qft, add, grover, period_find, shor" << endl;
+             cerr << "Usage examples:" << endl;
+             cerr << "  ./test                 (runs all tests)" << endl;
+             cerr << "  ./test shor 21" << endl;
+             cerr << "  ./test grover 6 3" << endl;
+             cerr << "  ./test period_find 7 15" << endl;
+            return 1;
+        }
+    } catch (const std::invalid_argument& ia) {
+        cerr << "Error: Invalid numerical argument provided for command '" << command << "'." << endl;
+        return 1;
+    } catch (const std::out_of_range& oor) {
+         cerr << "Error: Numerical argument out of range for command '" << command << "'." << endl;
+         return 1;
+	}
+
+    return 0; // Success
 }