diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index f7f2cad..7f6266f 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -7,7 +7,7 @@ jobs:
       matrix:
         os: ["ubuntu", "macos"]
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
       - run: rustup update
       - run: cargo build --verbose
       - run: cargo test --verbose
diff --git a/test/test.c b/test/test.c
index 354f624..45c6182 100644
--- a/test/test.c
+++ b/test/test.c
@@ -1,10 +1,28 @@
 #include "../tiktoken.h"
 #include <stdio.h>
+#include <stdlib.h>
 #include <string.h>
+#include <unistd.h>
 #include <sys/types.h>
 
-int main(void)
+int main(int argc, char *argv[])
 {
+    char *model = "gpt-4"; // default model
+    int opt;
+
+    while ((opt = getopt(argc, argv, "m:")) != -1)
+    {
+        switch (opt)
+        {
+        case 'm':
+            model = optarg;
+            break;
+        default:
+            fprintf(stderr, "Usage: %s [-m model]\n", argv[0]);
+            exit(EXIT_FAILURE);
+        }
+    }
+
     char *text = NULL;
     size_t len = 0;
     ssize_t nread;
@@ -22,7 +40,7 @@ int main(void)
         text[--nread] = '\0';
     }
 
-    CoreBPE *bpe = c_get_bpe_from_model("gpt-4");
+    CoreBPE *bpe = c_get_bpe_from_model(model);
     size_t n;
     size_t *tokens = c_corebpe_encode_with_special_tokens(bpe, text, &n);
 
@@ -39,4 +57,4 @@ int main(void)
     free(tokens);
 
     return 0;
-}
+}
\ No newline at end of file
diff --git a/test/test.sh b/test/test.sh
index 177e1d4..1fe2a04 100755
--- a/test/test.sh
+++ b/test/test.sh
@@ -5,16 +5,44 @@ export LD_LIBRARY_PATH="../target/debug"
 
 cc test.c -L ../target/debug/ -ltiktoken_c -o test
 
-OUTPUT=$(echo "I am a cat." | ./test) 
+# Initialize test result
+ALL_TESTS_PASSED=true
 
-EXPECTED="40 1097 264 8415 13"
+# Test with gpt-4 model
+OUTPUT_GPT4=$(echo "I am a cat." | ./test -m "gpt-4")
+EXPECTED_GPT4="40 1097 264 8415 13"
 
-echo "Expected: $EXPECTED"
-echo "     Got: $OUTPUT"
+echo "Testing gpt-4 model"
+echo "Expected: $EXPECTED_GPT4"
+echo "     Got: $OUTPUT_GPT4"
 
-if [ "$OUTPUT" = "$EXPECTED" ]; then
-  echo "Test passed successfully"
+if [ "$OUTPUT_GPT4" = "$EXPECTED_GPT4" ]; then
+  echo "Test passed successfully for gpt-4"
 else
-  echo "Test failed :("
+  echo "Test failed for gpt-4 :("
+  ALL_TESTS_PASSED=false
+fi
+
+# Test with gpt-4o model
+OUTPUT_GPT4O=$(echo "I am a cat." | ./test -m "gpt-4o")
+EXPECTED_GPT4O="40 939 261 9059 13"
+
+echo "Testing gpt-4o model"
+echo "Expected: $EXPECTED_GPT4O"
+echo "     Got: $OUTPUT_GPT4O"
+
+if [ "$OUTPUT_GPT4O" = "$EXPECTED_GPT4O" ]; then
+  echo "Test passed successfully for gpt-4o"
+else
+  echo "Test failed for gpt-4o :("
+  ALL_TESTS_PASSED=false
+fi
+
+# Final test result
+if [ "$ALL_TESTS_PASSED" = true ]; then
+  echo "All tests passed successfully"
+  exit 0
+else
+  echo "Some tests failed :("
   exit 1
 fi