add --skip and --only to test driver

controllers/pyctrl/driver.py

@@ -1,6 +1,7 @@
 import sys
 import os
 import re
+import argparse
 
 import pyaici.rest
 import pyaici.util
@@ -51,10 +52,38 @@ def main():
     js_mode = False
     cmt = "#"
 
-    files = sys.argv[1:]
-    if not files:
-        print("need some python files as input")
-        return
+    parser = argparse.ArgumentParser(
+        description="Run pyctrl or jsctrl tests",
+        prog="ctrldriver",
+    )
+
+    parser.add_argument(
+        "--skip",
+        "-s",
+        type=str,
+        default=[],
+        action="append",
+        help="skip tests matching string",
+    )
+
+    parser.add_argument(
+        "--only",
+        "-k",
+        type=str,
+        default=[],
+        action="append",
+        help="only run tests matching string",
+    )
+
+    parser.add_argument(
+        "test_file",
+        nargs="+",
+        help="files to test",
+    )
+
+    args = parser.parse_args()
+
+    files = args.test_file
 
     if files[0].endswith(".js"):
         js_mode = True
@@ -77,6 +106,10 @@ def main():
         else:
             tests = re.findall(r"^async def (test_\w+)\(.*", arg, flags=re.MULTILINE)
         for t in tests:
+            if any([s in t for s in args.skip]):
+                continue
+            if args.only and not any([s in t for s in args.only]):
+                continue
             if js_mode:
                 arg_t = f"{arg}\ntest({t});\n"
             else:

controllers/pyctrl/samples/test.py

@@ -49,10 +49,11 @@ async def test_hello():
     prompt = await aici.GetPrompt()
     print("prompt", prompt)
     await aici.gen_tokens(regex=r"[A-Z].*", max_tokens=5)
+    await aici.FixedTokens("\n2 +")
     l = aici.Label()
-    await aici.FixedTokens("\n2 + 2 = ")
+    await aici.FixedTokens(" 2 = ")
     await aici.gen_tokens(regex=r"\d+", max_tokens=1)
-    await aici.FixedTokens("\n3 + 3 = ", following=l)
+    await aici.FixedTokens(" 3 = ", following=l)
     await aici.gen_tokens(regex=r"\d+", max_tokens=1)
 
 
@@ -187,4 +188,4 @@ async def test_joke():
     await aici.gen_text(max_tokens=15)
 
 
-aici.test(test_hello())
+aici.test(test_drugs())

scripts/test-jsctrl.sh

@@ -7,4 +7,4 @@ cd $HERE/../controllers/jsctrl
 tsc --version || npm install -g typescript
 tsc -p samples
 PYTHONPATH=$HERE/../py \
-python3 ../pyctrl/driver.py samples/dist/test.js
+python3 ../pyctrl/driver.py samples/dist/test.js "$@"

scripts/test-pyctrl.sh

@@ -5,4 +5,4 @@ cd `dirname $0`
 HERE=`pwd`
 cd $HERE/../controllers/pyctrl
 PYTHONPATH=$HERE/../py \
-python3 driver.py samples/test*.py
+python3 driver.py samples/test*.py "$@"