2 files changed, 18 insertions, 6 deletions
diff --git a/compile.c b/compile.c
index 2b4c39e4..3170369d 100644
--- a/compile.c
+++ b/compile.c
@@ -306,11 +306,23 @@ CORD compile(ast_t *ast)
     case DocTest: {
         auto test = Match(ast, DocTest);
         CORD src = heap_strn(test->expr->start, (size_t)(test->expr->end - test->expr->start));
-        return CORD_asprintf(
-            "__test(%r, %r, %r);\n",
-            compile(WrapAST(test->expr, StringLiteral, .cord=src)),
-            compile(test->expr),
-            compile(WrapAST(test->expr, StringLiteral, .cord=test->output)));
+        if (test->expr->tag == Declare) {
+            auto decl = Match(test->expr, Declare);
+            return CORD_asprintf(
+                "__declare(%r, %r);\n__test(%r, %r, %r);\n",
+                compile(decl->var), compile(decl->value),
+                compile(WrapAST(test->expr, StringLiteral, .cord=src)),
+                compile(decl->var),
+                compile(WrapAST(test->expr, StringLiteral, .cord=test->output)));
+        } else if (test->expr->tag == Assign) {
+            errx(1, "Not implemented");
+        } else {
+            return CORD_asprintf(
+                "__test(%r, %r, %r);\n",
+                compile(WrapAST(test->expr, StringLiteral, .cord=src)),
+                compile(test->expr),
+                compile(WrapAST(test->expr, StringLiteral, .cord=test->output)));
+        }
     }
     case FieldAccess: {
         auto f = Match(ast, FieldAccess);
diff --git a/nextlang.h b/nextlang.h
index fda796c8..7d63659c 100644
--- a/nextlang.h
+++ b/nextlang.h
@@ -78,7 +78,7 @@ CORD as_cord(void *x, const char *fmt, ...);
 #define say(str) puts(CORD_to_const_char_star(__cord(str)))
 #define __test(src, expr, expected) do { \
         CORD __result = __cord(expr); \
-        say(CORD_catn(5, USE_COLOR ? "\x1b[33;1m>>>\x1b[0m " : ">>> ", src, USE_COLOR ? "\n\x1b[0;2m=\x1b[0;35m " : "\n=== ", __result, "\x1b[m")); \
+        say(CORD_catn(5, USE_COLOR ? "\x1b[33;1m>>>\x1b[0m " : ">>> ", src, USE_COLOR ? "\n\x1b[0;2m=\x1b[0;35m " : "\n= ", __result, "\x1b[m")); \
         if (expected && CORD_cmp(__result, expected)) { \
             errx(1, "I expected:\n%s but got:\n%s", CORD_to_const_char_star(expected), CORD_to_const_char_star(__result)); \
         } \