From 870e41912ac4c3bea803d1038bd17b3d66a07cd3 Mon Sep 17 00:00:00 2001
From: Shreyas Goenka <shreyas.goenka@databricks.com>
Date: Tue, 27 Aug 2024 11:42:33 +0200
Subject: [PATCH] from_type method is mostly complete

---
 libs/jsonschema/from_type.go      |  6 ++-
 libs/jsonschema/from_type_test.go | 68 +++++++++++++++++++++++--------
 2 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/libs/jsonschema/from_type.go b/libs/jsonschema/from_type.go
index 389c16695..d2bbb4550 100644
--- a/libs/jsonschema/from_type.go
+++ b/libs/jsonschema/from_type.go
@@ -91,8 +91,8 @@ func FromType(typ reflect.Type, fns []func(typ reflect.Type, s Schema) Schema) (
 		return Schema{}, err
 	}
 
-	for k, v := range c.definitions {
-		for _, fn := range fns {
+	for _, fn := range fns {
+		for k, v := range c.definitions {
 			c.definitions[k] = fn(c.seen[k], v)
 		}
 	}
@@ -134,6 +134,8 @@ func typePath(typ reflect.Type) string {
 	case typ.Name() != "":
 		return typ.Name()
 	default:
+		// Invariant. This function should return a non-empty string
+		// for all types.
 		panic("unexpected empty type name for type: " + typ.String())
 	}
 }
diff --git a/libs/jsonschema/from_type_test.go b/libs/jsonschema/from_type_test.go
index e8bdd5eba..29d007901 100644
--- a/libs/jsonschema/from_type_test.go
+++ b/libs/jsonschema/from_type_test.go
@@ -127,15 +127,6 @@ func TestFromTypeBasic(t *testing.T) {
 			s, err := FromType(tc.typ, nil)
 			assert.NoError(t, err)
 			assert.Equal(t, tc.expected, s)
-
-			// jsonSchema, err := json.MarshalIndent(s, "		", "	")
-			// assert.NoError(t, err)
-
-			// expectedJson, err := json.MarshalIndent(tc.expected, "		", "	")
-			// assert.NoError(t, err)
-
-			// t.Log("[DEBUG] actual: ", string(jsonSchema))
-			// t.Log("[DEBUG] expected: ", string(expectedJson))
 		})
 	}
 }
@@ -158,8 +149,6 @@ func TestGetStructFields(t *testing.T) {
 	assert.Equal(t, "B", fields[2].Name)
 }
 
-// TODO: Add other case coverage for all the tests below
-
 func TestFromTypeNested(t *testing.T) {
 	type Inner struct {
 		S string `json:"s"`
@@ -352,16 +341,65 @@ func TestFromTypeSelfReferential(t *testing.T) {
 	assert.Equal(t, expected, s)
 }
 
-// TODO: Add coverage for all errors returned by FromType.
 func TestFromTypeError(t *testing.T) {
+	// Maps with non-string keys should panic.
 	type mapOfInts map[int]int
-
 	assert.PanicsWithValue(t, "found map with non-string key: int", func() {
 		FromType(reflect.TypeOf(mapOfInts{}), nil)
 	})
+
+	// Unsupported types should return an error.
+	_, err := FromType(reflect.TypeOf(complex64(0)), nil)
+	assert.EqualError(t, err, "unsupported type: complex64")
 }
 
-// TODO: Add test that the fn argument ot from_type works as expected.
+func TestFromTypeFunctionsArg(t *testing.T) {
+	type myStruct struct {
+		S string `json:"s"`
+	}
+
+	strRef := "#/$defs/string"
+	expected := Schema{
+		Type: "object",
+		Definitions: map[string]any{
+			"string": Schema{
+				Type:        "string",
+				Description: "a string",
+				Enum:        []any{"a", "b", "c"},
+			},
+		},
+		Properties: map[string]*Schema{
+			"s": {
+				Reference: &strRef,
+			},
+		},
+		AdditionalProperties: false,
+		Required:             []string{"s"},
+	}
+
+	addDescription := func(typ reflect.Type, s Schema) Schema {
+		if typ.Kind() != reflect.String {
+			return s
+		}
+		s.Description = "a string"
+		return s
+	}
+
+	addEnums := func(typ reflect.Type, s Schema) Schema {
+		if typ.Kind() != reflect.String {
+			return s
+		}
+		s.Enum = []any{"a", "b", "c"}
+		return s
+	}
+
+	s, err := FromType(reflect.TypeOf(myStruct{}), []func(reflect.Type, Schema) Schema{
+		addDescription,
+		addEnums,
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, expected, s)
+}
 
 func TestTypePath(t *testing.T) {
 	type myStruct struct{}
@@ -412,8 +450,6 @@ func TestTypePath(t *testing.T) {
 		},
 	}
 
-	// TODO: support arbitary depth of maps and slices. Also add validation
-	// in this function that non-string keys are not allowed.
 	for _, tc := range tcases {
 		t.Run(tc.typ.String(), func(t *testing.T) {
 			assert.Equal(t, tc.path, typePath(tc.typ))