From be68241d8e1a9a5be73f1cd072f753701e13cab9 Mon Sep 17 00:00:00 2001
From: Samir Patel <48686912+54mir@users.noreply.github.com>
Date: Wed, 9 Feb 2022 12:33:12 -0600
Subject: [PATCH] add test

---
 pql/ast_test.go | 104 ++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 104 insertions(+)

diff --git a/pql/ast_test.go b/pql/ast_test.go
index b095fe878..d77fe623b 100644
--- a/pql/ast_test.go
+++ b/pql/ast_test.go
@@ -2,6 +2,7 @@
 package pql_test
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/molecula/featurebase/v3/pql"
@@ -52,3 +53,106 @@ func TestCondition_StringWithSubj(t *testing.T) {
 		}
 	}
 }
+
+func TestQuery_ExpandVars(t *testing.T) {
+	tests := []struct {
+		name    string
+		input   string
+		output  string
+		vars    map[string]interface{}
+		wantErr bool
+	}{
+		{
+			name:   "ExpandRowEQInterior",
+			input:  `count(row(animal=$var1))`,
+			output: `Count(Union(Row(animal="cat"), Row(animal="dog"), Row(animal="pig")))`,
+			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog", "pig"}},
+		},
+		{
+			name:   "ExpandRowEQExterior",
+			input:  `row(animal=$var1)`,
+			output: `Union(Row(animal="cat"))`,
+			vars:   map[string]interface{}{"var1": []interface{}{"cat"}},
+		},
+		{
+			name:   "ExpandRowGT",
+			input:  `count(row(num>$var1))`,
+			output: `Count(Union(Row(num>5), Row(num>10)))`,
+			vars:   map[string]interface{}{"var1": []interface{}{5, 10}},
+		},
+		{
+			name:   "ExpandRowNOT",
+			input:  `count(row(num!=$var1))`,
+			output: `Count(Union(Row(num!=5), Row(num!=10)))`,
+			vars:   map[string]interface{}{"var1": []interface{}{5, 10}},
+		},
+		{
+			name:   "ExpandRowLTString",
+			input:  `count(row(num<$var1))`,
+			output: `Count(Union(Row(num<"cat"), Row(num<"dog")))`,
+			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog"}},
+		},
+		{
+			name:   "ExpandRowsInterior",
+			input:  `GroupBy(rows($var1), limit=5)`,
+			output: `GroupBy(Rows(_field="cat"), Rows(_field="dog"), limit=5)`,
+			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog"}},
+		},
+		{
+			name:   "ExpandRowsExterior",
+			input:  `rows($var1)`,
+			output: `Rows(_field="cat")` + "\n" + `Rows(_field="dog")`,
+			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog"}},
+		},
+		{
+			name:   "ExpandRowAndRows",
+			input:  `GroupBy(Rows($animal), limit=7, filter=Row(size=$size))`,
+			output: `GroupBy(Rows(_field="cat"), Rows(_field="dog"), filter=Union(Row(size="lg"), Row(size="md")), limit=7)`,
+			vars:   map[string]interface{}{"animal": []interface{}{"cat", "dog"}, "size": []interface{}{"lg", "md"}},
+		},
+		{
+			name:    "ExpandBad",
+			input:   `$animal`,
+			vars:    map[string]interface{}{"animal": []interface{}{"cat", "dog"}, "columns": []interface{}{5, 10}},
+			wantErr: true,
+		},
+		{
+			name:    "ExpandBad2",
+			input:   `GroupBy($animal)`,
+			output:  `Intersect(ConstRow(columns=[5, 10]), Union(Row(animal="cat"), Row(animal="dog")))`,
+			wantErr: true,
+		},
+		{
+			name:   "ExpandAsCSV",
+			input:  `Intersect(ConstRow(columns=$var2), Row(animal=$var1))`,
+			output: `Intersect(ConstRow(columns=[5,10]), Union(Row(animal="cat"), Row(animal="dog")))`,
+			vars:   map[string]interface{}{"var1": []interface{}{"cat", "dog"}, "var2": []interface{}{5, 10}},
+		},
+		{
+			name:   "ExpandPercentile",
+			input:  `Percentile(field="bytes", nth=99.0, filter=Row(level=$animal))`,
+			output: `Percentile(field="bytes", filter=Union(Row(level="cat"), Row(level="dog")), nth=99)`,
+			vars:   map[string]interface{}{"animal": []interface{}{"cat", "dog"}, "columns": []interface{}{5, 10}},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			q, err := pql.NewParser(strings.NewReader(tt.input)).Parse()
+			if err != nil {
+				if !tt.wantErr {
+					t.Errorf("Parse error = %v, wantErr %v", err, tt.wantErr)
+				}
+				return
+			}
+
+			got, err := q.ExpandVars(tt.vars)
+			if err != nil {
+				t.Errorf("Query.ExpandVars() error = %v", err)
+				return
+			}
+			if tt.output != got.String() {
+				t.Errorf("got %v, want %v", got, tt.output)
+			}
+		})
+	}
+}