From 8c215954a55cfb5caa55bcd86f2745a8544c3ad8 Mon Sep 17 00:00:00 2001
From: Kyle Ellrott <kellrott@gmail.com>
Date: Thu, 19 Dec 2019 12:17:41 -0800
Subject: [PATCH] Adding some unit tests to check of correct identification of
 Distinct method data usage

---
 engine/inspect/inspect.go           |  3 +++
 engine/inspect/test/inspect_test.go | 38 +++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/engine/inspect/inspect.go b/engine/inspect/inspect.go
index 5a154334..0204d551 100644
--- a/engine/inspect/inspect.go
+++ b/engine/inspect/inspect.go
@@ -99,6 +99,9 @@ func PipelineStepOutputs(stmts []*gripql.GraphStatement) map[string][]string {
 			fields := protoutil.AsStringList(gs.GetDistinct())
 			for _, f := range fields {
 				n := jsonpath.GetNamespace(f)
+				if n == "__current__" {
+					out[steps[i]] = []string{"*"}
+				}
 				if a, ok := asMap[n]; ok {
 					out[a] = []string{"*"}
 				}
diff --git a/engine/inspect/test/inspect_test.go b/engine/inspect/test/inspect_test.go
index b1ab0fb4..ddac77d1 100644
--- a/engine/inspect/test/inspect_test.go
+++ b/engine/inspect/test/inspect_test.go
@@ -156,3 +156,41 @@ func TestPathFind(t *testing.T) {
 		}
 	}
 }
+
+func TestDistinct(t *testing.T) {
+	q := gripql.NewQuery()
+	o := q.V().HasLabel("Person").As("person").Out("friend").Distinct("$person.name", "$.name").Count()
+	out := inspect.PipelineStepOutputs(o.Statements)
+	fmt.Printf("%#v\n", out)
+
+	if x, ok := out["2"]; ok {
+		if !setcmp.ContainsString(x, "*") {
+			t.Errorf("Required data from step 2 not recognized")
+		}
+	} else {
+		t.Errorf("Data load for step 2 not recognized")
+	}
+
+	o = q.V().HasLabel("Person").As("person").Out("friend").Distinct("$.name").Count()
+	out = inspect.PipelineStepOutputs(o.Statements)
+	fmt.Printf("%#v\n", out)
+	if x, ok := out["2"]; ok {
+		if !setcmp.ContainsString(x, "*") {
+			t.Errorf("Required data from step 2 not recognized")
+		}
+	} else {
+		t.Errorf("Data load for step 2 not recognized")
+	}
+
+	o = q.V().HasLabel("Person").As("person").Out("friend").Distinct("$person.name").Out("friend").Distinct("$.name").Count()
+	out = inspect.PipelineStepOutputs(o.Statements)
+	fmt.Printf("%#v\n", out)
+	if x, ok := out["3"]; ok {
+		if !setcmp.ContainsString(x, "*") {
+			t.Errorf("Required data from step 3 not recognized")
+		}
+	} else {
+		t.Errorf("Data load for step 3 not recognized")
+	}
+
+}