From 274b5fef18c7a63840b75c17d7a229de49862e24 Mon Sep 17 00:00:00 2001
From: Matthew Davidson <matthew@modulolotus.net>
Date: Mon, 8 Apr 2024 20:05:25 +0700
Subject: [PATCH] test: Add GENERATE * EXCEPT tests

---
 test/inferenceql/query/base/parser_test.cljc       | 3 +++
 test/inferenceql/query/permissive/parser_test.cljc | 6 ++++++
 test/inferenceql/query/permissive_test.cljc        | 5 ++++-
 test/inferenceql/query/plan_test.cljc              | 3 ++-
 test/inferenceql/query/strict/parser_test.cljc     | 6 ++++++
 test/inferenceql/query/strict_test.cljc            | 7 +++++--
 6 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/test/inferenceql/query/base/parser_test.cljc b/test/inferenceql/query/base/parser_test.cljc
index 7a2fdc5..2d0d533 100644
--- a/test/inferenceql/query/base/parser_test.cljc
+++ b/test/inferenceql/query/base/parser_test.cljc
@@ -107,3 +107,6 @@
 
 (deftest parse-generative-join-success
   (parses "table GENERATIVE JOIN model" :generative-join-expr))
+
+(deftest parse-generate-star
+  (parses "GENERATE * UNDER model" :generate-expr))
diff --git a/test/inferenceql/query/permissive/parser_test.cljc b/test/inferenceql/query/permissive/parser_test.cljc
index f65a3c2..a353ce4 100644
--- a/test/inferenceql/query/permissive/parser_test.cljc
+++ b/test/inferenceql/query/permissive/parser_test.cljc
@@ -35,3 +35,9 @@
     "data GENERATIVE JOIN model CONDITIONED BY x = 0"
     "data GENERATIVE JOIN model CONSTRAINED BY x > 0"
     "data GENERATIVE JOIN model GIVEN x = 0"))
+
+(deftest generate-valid
+  (are [s] (not (insta/failure? (parser/parse s :start :generate-expr)))
+    "GENERATE * UNDER model"
+    "GENERATE foo, bar UNDER model"
+    "GENERATE * EXCEPT (foo, bar) UNDER model"))
diff --git a/test/inferenceql/query/permissive_test.cljc b/test/inferenceql/query/permissive_test.cljc
index 5ea2f8d..5f92b41 100644
--- a/test/inferenceql/query/permissive_test.cljc
+++ b/test/inferenceql/query/permissive_test.cljc
@@ -141,7 +141,10 @@
     "GENERATE VAR x, VAR y UNDER model"
 
     "GENERATE x, y, z UNDER model"
-    "GENERATE VAR x, VAR y, VAR z UNDER model"))
+    "GENERATE VAR x, VAR y, VAR z UNDER model"
+
+    "GENERATE * EXCEPT (x, y) UNDER model"
+    "GENERATE * EXCEPT (VAR x, VAR y) UNDER model"))
 
 (deftest approximate-mutual-info
   (are [permissive strict] (= (strict.parser/parse strict :start :scalar-expr)
diff --git a/test/inferenceql/query/plan_test.cljc b/test/inferenceql/query/plan_test.cljc
index 19cc6cb..772a348 100644
--- a/test/inferenceql/query/plan_test.cljc
+++ b/test/inferenceql/query/plan_test.cljc
@@ -232,7 +232,8 @@
       "GENERATE VAR x UNDER model" ["x"]
       "GENERATE VAR x, VAR y UNDER model" ["x" "y"]
       "GENERATE VAR y, VAR x UNDER model" ["y" "x"]
-      "GENERATE VAR x, VAR y, VAR z UNDER model" ["x" "y" "z"]))
+      "GENERATE * EXCEPT (VAR x) UNDER model" ["y"]
+      "GENERATE * EXCEPT (VAR x, VAR y) UNDER model" nil))
   (testing "values"
     (let [rel (eval "GENERATE VAR x UNDER model" {"model" model})]
       (doseq [tup (relation/tuples (take 1 #_ 5 rel))]
diff --git a/test/inferenceql/query/strict/parser_test.cljc b/test/inferenceql/query/strict/parser_test.cljc
index 9226d73..7054e4c 100644
--- a/test/inferenceql/query/strict/parser_test.cljc
+++ b/test/inferenceql/query/strict/parser_test.cljc
@@ -42,6 +42,12 @@
     "data GENERATIVE JOIN model CONDITIONED BY VAR x = 0"
     "data GENERATIVE JOIN model CONSTRAINED BY VAR x > 0"))
 
+(deftest generate-valid
+  (are [s] (not (insta/failure? (parser/parse s :start :generate-expr)))
+    "GENERATE * UNDER model"
+    "GENERATE VAR foo, VAR bar UNDER model"
+    "GENERATE * EXCEPT (VAR foo, VAR bar) UNDER model"))
+
 (deftest conditioned-by-valid
   (are [s] (not (insta/failure? (parser/parse s)))
     "SELECT * FROM (GENERATE * UNDER model CONDITIONED BY VAR x = x)"
diff --git a/test/inferenceql/query/strict_test.cljc b/test/inferenceql/query/strict_test.cljc
index c15ec79..16b0c1f 100644
--- a/test/inferenceql/query/strict_test.cljc
+++ b/test/inferenceql/query/strict_test.cljc
@@ -286,14 +286,17 @@
 
 ;;; Generate
 
-#_
+
 (deftest generate-generates-correct-columns
-  (testing "generate"
+  (testing "GENERATE"
     (let [model simple-model
           q #(q % [] {"model" model})]
       (testing "with star "
         (doseq [result (q "SELECT * FROM (GENERATE * UNDER model) LIMIT 10")]
           (is (= #{"x" "y"} (set (keys result))))))
+      (testing "with star except"
+        (doseq [result (q "SELECT * FROM (GENERATE * EXCEPT (VAR x) UNDER model) LIMIT 10")]
+          (is (= #{"y"} (set (keys result))))))
       (testing "with a single variable"
         (doseq [result (q "SELECT * FROM (GENERATE VAR y UNDER model) LIMIT 10")]
           (is (= #{"y"} (set (keys result))))))