Improved C test

test/test.c

@@ -1,23 +1,42 @@
 #include "../tiktoken.h"
 #include <stdio.h>
+#include <string.h>
+#include <sys/types.h>
 
-int main(int argc, char *argv[]){
-    if (argc < 2) {
-        printf("Error: No text argument provided.\n");
+int main(void)
+{
+    char *text = NULL;
+    size_t len = 0;
+    ssize_t nread;
+
+    nread = getline(&text, &len, stdin);
+    if (nread == -1)
+    {
+        printf("Error: No text received from stdin or reading error.\n");
         return 1;
     }
 
-    char *text = argv[1];
-    printf("%s\n", text);
+    // getline retains newline character. We may want to remove this.
+    if (nread > 0 && text[nread - 1] == '\n')
+    {
+        text[--nread] = '\0';
+    }
 
     CoreBPE *bpe = c_get_bpe_from_model("gpt-4");
     size_t n;
     size_t *tokens = c_corebpe_encode_with_special_tokens(bpe, text, &n);
 
-    for(size_t i=0; i<n; i++){
-        printf("%zu ", tokens[i]);
+    for (size_t i = 0; i < n; i++)
+    {
+        printf("%zu", tokens[i]);
+        if (i < n - 1)
+        {
+            printf(" ");
+        }
     }
-    printf("\n");
+
+    free(text); // do not forget to free allocated memory
+    free(tokens);
 
     return 0;
 }

test/test.sh

@@ -1,3 +1,20 @@
 #!/bin/sh
+
+# Add missing import
+export LD_LIBRARY_PATH="../target/debug"
+
 cc test.c -L ../target/debug/ -ltiktoken_c -o test
-LD_LIBRARY_PATH="../target/debug" ./test "I am a cat."
+
+OUTPUT=$(echo "I am a cat." | ./test) 
+
+EXPECTED="40 1097 264 8415 13"
+
+echo "Expected: $EXPECTED"
+echo "     Got: $OUTPUT"
+
+if [ "$OUTPUT" = "$EXPECTED" ]; then
+  echo "Test passed successfully"
+else
+  echo "Test failed :("
+  exit 1
+fi