From a0bad4d6e955f824624a16984d545fdae7c455c5 Mon Sep 17 00:00:00 2001
From: Muyao CHEN <muyao.chen@intersec.com>
Date: Fri, 20 Sep 2024 21:04:25 +0200
Subject: [PATCH] walk: handle functions...

---
 walk/walk.go      |  5 +++++
 walk/walk_test.go | 15 +++++++++++++++
 2 files changed, 20 insertions(+)

diff --git a/walk/walk.go b/walk/walk.go
index 871f1bf..acb72af 100644
--- a/walk/walk.go
+++ b/walk/walk.go
@@ -35,6 +35,11 @@ func Walk(x interface{}, fn func(string)) {
 				break
 			}
 		}
+	case reflect.Func:
+		valFnResult := val.Call(nil)
+		for _, res := range valFnResult {
+			walkValue(res)
+		}
 	}
 }
 
diff --git a/walk/walk_test.go b/walk/walk_test.go
index d24bf7c..40fb741 100644
--- a/walk/walk_test.go
+++ b/walk/walk_test.go
@@ -115,6 +115,21 @@ func TestWalk(t *testing.T) {
 			t.Errorf("got %v want %v", got, want)
 		}
 	})
+
+	t.Run("with functions", func(t *testing.T) {
+		aFunction := func() (Profile, Profile) {
+			return Profile{21, "Berlin"}, Profile{28, "Beijing"}
+		}
+
+		var got []string
+		want := []string{"Berlin", "Beijing"}
+		Walk(aFunction, func(input string) {
+			got = append(got, input)
+		})
+		if !reflect.DeepEqual(got, want) {
+			t.Errorf("got %v want %v", got, want)
+		}
+	})
 }
 
 func assertContains(t testing.TB, haystack []string, needle string) {