From 35f94b3c48b047d286c8e08a7d18a5f0f87f9058 Mon Sep 17 00:00:00 2001
From: Ben Kraft <benkraft@khanacademy.org>
Date: Thu, 2 Apr 2020 18:16:03 -0700
Subject: [PATCH] some tests for type-generation

---
 generate/types.go      |   9 ++-
 generate/types_test.go | 137 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 145 insertions(+), 1 deletion(-)
 create mode 100644 generate/types_test.go

diff --git a/generate/types.go b/generate/types.go
index ac05a44..3614f64 100644
--- a/generate/types.go
+++ b/generate/types.go
@@ -32,6 +32,7 @@ func writeSelectionSetStruct(builder *strings.Builder, selectionSet ast.Selectio
 			if selection.Alias != "" {
 				jsonName = selection.Alias
 			} else {
+				// TODO: is this case needed? tests don't seem to get here.
 				jsonName = selection.Name
 			}
 			// We need an exportable name for JSON-marshaling.
@@ -40,6 +41,11 @@ func writeSelectionSetStruct(builder *strings.Builder, selectionSet ast.Selectio
 			builder.WriteString(goName)
 			builder.WriteRune(' ')
 
+			if selection.Definition == nil {
+				// Unclear why gqlparser hasn't already rejected this,
+				// but empirically it might not.
+				return fmt.Errorf("undefined selection %v", selection)
+			}
 			writeType(builder, selection.Definition.Type, selection.SelectionSet, schema)
 
 			if jsonName != goName {
@@ -77,7 +83,8 @@ func writeType(builder *strings.Builder, typ *ast.Type, selectionSet ast.Selecti
 		// Type is a list.
 		builder.WriteString("[]")
 		typ = typ.Elem
-	} else if !typ.NonNull { // no need for pointer if we have a list
+	}
+	if !typ.NonNull {
 		builder.WriteString("*")
 	}
 
diff --git a/generate/types_test.go b/generate/types_test.go
new file mode 100644
index 0000000..2e19457
--- /dev/null
+++ b/generate/types_test.go
@@ -0,0 +1,137 @@
+package generate
+
+import (
+	"go/format"
+	"testing"
+
+	"github.com/vektah/gqlparser"
+	"github.com/vektah/gqlparser/ast"
+)
+
+func gofmt(src string) (string, error) {
+	formatted, err := format.Source([]byte(src))
+	if err != nil {
+		return src, err
+	}
+	return string(formatted), nil
+}
+
+func TestTypeForOperation(t *testing.T) {
+	schema, err := gqlparser.LoadSchema(&ast.Source{Name: "test schema", Input: `
+		type AuthMethod {
+			provider: String
+			email: String
+		}
+
+		type User {
+			id: ID!
+			name: String
+			emails: [String!]!
+			emailsOrNull: [String!]
+			emailsWithNulls: [String]!
+			emailsWithNullsOrNull: [String]
+			authMethods: [AuthMethod!]!
+		}
+
+		type Query {
+			user: User
+		}
+	`})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	tests := []struct {
+		name           string
+		operation      string
+		expectedGoType string
+	}{{
+		"SimpleQuery",
+		`{ user { id } }`,
+		`struct{
+			User *struct {
+				Id string ` + "`json:\"id\"`" + `
+			} ` + "`json:\"user\"`" + `
+		}`,
+	}, {
+		"QueryWithAlias",
+		`{ User: user { ID: id } }`,
+		`struct{
+			User *struct {
+				ID string
+			}
+		}`,
+		// Here on out, we use aliases, just because aliases are a lot less
+		// annoying to write in Go strings than Go struct tags.
+	}, {
+		"QueryWithSlices",
+		`{
+			User: user {
+				Emails: emails
+				EmailsOrNull: emailsOrNull
+				EmailsWithNulls: emailsWithNulls
+				EmailsWithNullsOrNull: emailsWithNullsOrNull
+			}
+		}`,
+		`struct{
+			User *struct {
+				Emails                []string
+				EmailsOrNull          []string
+				EmailsWithNulls       []*string
+				EmailsWithNullsOrNull []*string
+			}
+		}`,
+	}, {
+		"QueryWithStructs",
+		`{
+			User: user {
+				AuthMethods: authMethods {
+					Provider: provider
+					Email: email
+				}
+			}
+		}`,
+		`struct{
+			User *struct {
+				AuthMethods []struct {
+					Provider *string
+					Email    *string
+				}
+			}
+		}`,
+	}}
+
+	for _, test := range tests {
+		test := test
+		t.Run(test.name, func(t *testing.T) {
+			expectedGoType, err := gofmt("type Response " + test.expectedGoType)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			queryDoc, graphqlError := gqlparser.LoadQuery(schema, test.operation)
+			if graphqlError != nil {
+				t.Fatal(graphqlError)
+			}
+
+			if len(queryDoc.Operations) != 1 {
+				t.Fatalf("got %v operations, want 1", len(queryDoc.Operations))
+			}
+
+			goType, err := typeForOperation(queryDoc.Operations[0], schema)
+			if err != nil {
+				t.Error(err)
+			}
+
+			// gofmt before comparing.
+			goType, err = gofmt("type Response " + goType)
+			if err != nil {
+				t.Error(err)
+			}
+
+			if goType != expectedGoType {
+				t.Errorf("got:\n%v\nwant:\n%v\n", goType, expectedGoType)
+			}
+		})
+	}
+}