From c12d485ec27e7d6519ad4811cd695a6639ca4040 Mon Sep 17 00:00:00 2001
From: k1LoW <k1lowxb@gmail.com>
Date: Mon, 25 Sep 2023 09:38:17 +0900
Subject: [PATCH 1/2] Response() supports accepting any value

With this improvement, structures for responses generated by protoc-gen-go can be used as is.
---
 grpcstub.go      | 21 +++++++++++++++++++--
 grpcstub_test.go | 24 ++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/grpcstub.go b/grpcstub.go
index 970148a..c9e74de 100644
--- a/grpcstub.go
+++ b/grpcstub.go
@@ -128,6 +128,7 @@ type matcher struct {
 	matchFuncs []matchFunc
 	handler    handlerFunc
 	requests   []*Request
+	t          *testing.T
 	mu         sync.RWMutex
 }
 
@@ -283,6 +284,7 @@ func (s *Server) startServer() {
 func (s *Server) Match(fn func(r *Request) bool) *matcher {
 	m := &matcher{
 		matchFuncs: []matchFunc{fn},
+		t:          s.t,
 	}
 	s.mu.Lock()
 	defer s.mu.Unlock()
@@ -305,6 +307,7 @@ func (s *Server) Service(service string) *matcher {
 	fn := serviceMatchFunc(service)
 	m := &matcher{
 		matchFuncs: []matchFunc{fn},
+		t:          s.t,
 	}
 	s.matchers = append(s.matchers, m)
 	return m
@@ -336,6 +339,7 @@ func (s *Server) Method(method string) *matcher {
 	fn := methodMatchFunc(method)
 	m := &matcher{
 		matchFuncs: []matchFunc{fn},
+		t:          s.t,
 	}
 	s.matchers = append(s.matchers, m)
 	return m
@@ -400,7 +404,20 @@ func (m *matcher) Handler(fn func(r *Request) *Response) {
 }
 
 // Response set handler which return response.
-func (m *matcher) Response(message map[string]any) *matcher {
+func (m *matcher) Response(message any) *matcher {
+	mm := map[string]any{}
+	switch message.(type) {
+	case map[string]any:
+		mm = message.(map[string]any)
+	default:
+		b, err := json.Marshal(message)
+		if err != nil {
+			m.t.Fatalf("failed to convert message: %v", err)
+		}
+		if err := json.Unmarshal(b, &mm); err != nil {
+			m.t.Fatalf("failed to convert message: %v", err)
+		}
+	}
 	prev := m.handler
 	m.handler = func(r *Request, md protoreflect.MethodDescriptor) *Response {
 		var res *Response
@@ -409,7 +426,7 @@ func (m *matcher) Response(message map[string]any) *matcher {
 		} else {
 			res = prev(r, md)
 		}
-		res.Messages = append(res.Messages, message)
+		res.Messages = append(res.Messages, mm)
 		return res
 	}
 	return m
diff --git a/grpcstub_test.go b/grpcstub_test.go
index ffa8b5e..5c29802 100644
--- a/grpcstub_test.go
+++ b/grpcstub_test.go
@@ -656,3 +656,27 @@ func TestRequestStringer(t *testing.T) {
 		})
 	}
 }
+
+func TestResponseAny(t *testing.T) {
+	ctx := context.Background()
+	ts := NewServer(t, "testdata/route_guide.proto")
+	t.Cleanup(func() {
+		ts.Close()
+	})
+	ts.Service("routeguide.RouteGuide").Method("GetFeature").Response(&routeguide.Feature{
+		Name: "hello",
+	})
+
+	client := routeguide.NewRouteGuideClient(ts.Conn())
+	res, err := client.GetFeature(ctx, &routeguide.Point{
+		Latitude:  10,
+		Longitude: 13,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	got := res.Name
+	if want := "hello"; got != want {
+		t.Errorf("got %v\nwant %v", got, want)
+	}
+}

From 989c335e1e4e0e580d1b2586159f3a0932e5278d Mon Sep 17 00:00:00 2001
From: k1LoW <k1lowxb@gmail.com>
Date: Mon, 25 Sep 2023 09:46:49 +0900
Subject: [PATCH 2/2] Fix lint warn

---
 grpcstub.go | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/grpcstub.go b/grpcstub.go
index c9e74de..ca93bf4 100644
--- a/grpcstub.go
+++ b/grpcstub.go
@@ -406,11 +406,11 @@ func (m *matcher) Handler(fn func(r *Request) *Response) {
 // Response set handler which return response.
 func (m *matcher) Response(message any) *matcher {
 	mm := map[string]any{}
-	switch message.(type) {
+	switch v := message.(type) {
 	case map[string]any:
-		mm = message.(map[string]any)
+		mm = v
 	default:
-		b, err := json.Marshal(message)
+		b, err := json.Marshal(v)
 		if err != nil {
 			m.t.Fatalf("failed to convert message: %v", err)
 		}