Updates in response to code review feedback

Author: yolken-segment

pkg/util/template.go

@@ -18,12 +18,11 @@ import (
 
 var (
 	extraTemplateFuncs = template.FuncMap{
-		"lookup":    lookup,
-		"toYaml":    toYaml,
-		"urlEncode": url.QueryEscape,
+		"lookup":     lookup,
+		"pathLookup": pathLookup,
+		"toYaml":     toYaml,
+		"urlEncode":  url.QueryEscape,
 	}
-
-	zeroValue = reflect.Value{}
 )
 
 // ApplyTemplate runs golang templating on all files in the provided path,
@@ -220,36 +219,49 @@ func configMapEntriesGenerator(
 	}
 }
 
-// lookup does a dot-separated path lookup on the input struct or map. If the input or any of
-// its children on the targeted path are not a map or struct, it returns nil.
-func lookup(path string, input interface{}) interface{} {
+// lookup does a dot-separated path lookup on the input map. If a key on the path is
+// not found, it returns nil. If the input or any of its children on the lookup path is not
+// a map, it returns an error.
+func lookup(input interface{}, path string) (interface{}, error) {
 	obj := reflect.ValueOf(input)
-
 	components := strings.Split(path, ".")
 
 	for i := 0; i < len(components); {
-		component := components[i]
-
 		switch obj.Kind() {
-		case reflect.Struct:
-			obj = obj.FieldByName(component)
-			i++
 		case reflect.Map:
-			obj = obj.MapIndex(reflect.ValueOf(component))
+			obj = obj.MapIndex(reflect.ValueOf(components[i]))
 			i++
 		case reflect.Ptr, reflect.Interface:
+			if obj.IsNil() {
+				return nil, nil
+			}
+
 			// Get the thing being pointed to or interfaced, don't advance index
 			obj = obj.Elem()
 		default:
-			// Got an unexpected type
-			return nil
+			if obj.IsValid() {
+				// An object was found, but it's not a map. Return an error.
+				return nil, fmt.Errorf(
+					"Tried to traverse a value that's not a map (kind=%s)",
+					obj.Kind(),
+				)
+			}
+
+			// An intermediate key wasn't found
+			return nil, nil
 		}
 	}
 
-	if obj == zeroValue {
-		return nil
+	if !obj.IsValid() {
+		// The last key wasn't found
+		return nil, nil
 	}
-	return obj.Interface()
+	return obj.Interface(), nil
+}
+
+// pathLookup is the same as lookup, but with the arguments flipped.
+func pathLookup(path string, input interface{}) (interface{}, error) {
+	return lookup(input, path)
 }
 
 func toYaml(input interface{}) (string, error) {

pkg/util/template_test.go

@@ -137,44 +137,87 @@ func fileContents(t *testing.T, path string) string {
 	return string(contents)
 }
 
-type TestStruct struct {
-	Key    string
-	Inner  TestStructInner
-	Inner2 *TestStructInner
-}
-type TestStructInner struct {
-	Map map[string]interface{}
-}
-
 func TestLookup(t *testing.T) {
-	s := TestStruct{
-		Key: "value0",
-		Inner: TestStructInner{
-			Map: map[string]interface{}{
-				"key1": "value1",
-				"key2": map[string]interface{}{
-					"key3": map[string]interface{}{
-						"key4": "value4",
-					},
-					"key5": 1234,
-				},
+	m := map[string]interface{}{
+		"key1": "value1",
+		"key2": map[string]interface{}{
+			"key3": map[string]interface{}{
+				"key4": "value4",
 			},
+			"key5": 1234,
 		},
-		Inner2: &TestStructInner{
-			Map: map[string]interface{}{
-				"key6": "value6",
-			},
+		"key6": nil,
+	}
+
+	type testCase struct {
+		input          interface{}
+		path           string
+		expectedResult interface{}
+		expectErr      bool
+	}
+
+	testCases := []testCase{
+		{
+			input:          m,
+			path:           "bad-key",
+			expectedResult: nil,
+		},
+		{
+			input:          m,
+			path:           "",
+			expectedResult: nil,
+		},
+		{
+			input:          nil,
+			path:           "key1",
+			expectedResult: nil,
 		},
+		{
+			input:          "not a map",
+			path:           "key1",
+			expectedResult: nil,
+			expectErr:      true,
+		},
+		{
+			input:          m,
+			path:           "key1",
+			expectedResult: "value1",
+		},
+		{
+			input:          &m,
+			path:           "key1",
+			expectedResult: "value1",
+		},
+		{
+			input:          m,
+			path:           "key1.not-a-map",
+			expectedResult: nil,
+			expectErr:      true,
+		},
+		{
+			input:          m,
+			path:           "key2.key3.key4",
+			expectedResult: "value4",
+		},
+		{
+			input:          m,
+			path:           "key2.key5",
+			expectedResult: 1234,
+		},
+		{
+			input:          m,
+			path:           "key6.nil-key",
+			expectedResult: nil,
+		},
+	}
+
+	for index, tc := range testCases {
+		result, err := lookup(tc.input, tc.path)
+		assert.Equal(t, tc.expectedResult, result, "Unexpected result for case %d", index)
+		if tc.expectErr {
+			assert.Error(t, err, "Did not get expected error in case %d", index)
+		} else {
+			assert.NoError(t, err, "Got unexpected error in case %d", index)
+		}
 	}
-	assert.Equal(t, "value0", lookup("Key", s))
-	assert.Equal(t, nil, lookup("bad-key", s))
-	assert.Equal(t, nil, lookup("", s))
-	assert.Equal(t, nil, lookup("key", "not a map"))
-	assert.Equal(t, "value1", lookup("Inner.Map.key1", s))
-	assert.Equal(t, "value1", lookup("Inner.Map.key1", &s))
-	assert.Equal(t, "value4", lookup("Inner.Map.key2.key3.key4", s))
-	assert.Equal(t, 1234, lookup("Inner.Map.key2.key5", s))
-	assert.Equal(t, nil, lookup("Inner.Map.non-existent-key", s))
-	assert.Equal(t, nil, lookup("Inner.Map.key2.non-existent-key", s))
-	assert.Equal(t, "value6", lookup("Inner2.Map.key6", s))
 }