1 files changed, 35 insertions, 18 deletions

diff --git a/gemini/response_test.go b/gemini/response_test.go
index 7ffb585..616fac4 100644
--- a/gemini/response_test.go
+++ b/gemini/response_test.go
@@ -248,7 +248,7 @@ func TestResponseClose(t *testing.T) {
 
 	resp = &gemini.Response{
 		Status: gemini.StatusInput,
-		Meta: "give me more",
+		Meta:   "give me more",
 	}
 
 	if err := resp.Close(); err != nil {
@@ -272,7 +272,7 @@ func TestResponseWriteTo(t *testing.T) {
 	clone := func(resp *gemini.Response) *gemini.Response {
 		other := &gemini.Response{
 			Status: resp.Status,
-			Meta: resp.Meta,
+			Meta:   resp.Meta,
 		}
 
 		if resp.Body != nil {
@@ -291,27 +291,44 @@ func TestResponseWriteTo(t *testing.T) {
 			other.Body = bytes.NewBuffer(buf2)
 		}
 
-		return resp
+		return other
 	}
 
-	r1 := &gemini.Response{
-		Status: gemini.StatusSuccess,
-		Meta: "text/gemini",
-		Body: bytes.NewBufferString("the body goes here"),
+	table := []struct {
+		name     string
+		response *gemini.Response
+	}{
+		{
+			name: "simple success",
+			response: gemini.Success(
+				"text/gemini",
+				bytes.NewBufferString("the body goes here"),
+			),
+		},
+		{
+			name:     "no body",
+			response: gemini.Input("need more pls"),
+		},
 	}
-	r2 := clone(r1)
 
-	wtbuf := &bytes.Buffer{}
-	if _, err := r1.WriteTo(wtbuf); err != nil {
-		t.Fatalf("response.WriteTo(): %s", err.Error())
-	}
+	for _, test := range table {
+		t.Run(test.name, func(t *testing.T) {
+			r1 := test.response
+			r2 := clone(test.response)
 
-	rdbuf := make([]byte, wtbuf.Len())
-	if n, err := r2.Read(rdbuf); err != nil {
-		t.Fatalf("response.Read() -> %d: %s", n, err.Error())
-	}
+			rdbuf, err := io.ReadAll(r1)
+			if err != nil {
+				t.Fatalf("response.Read(): %s", err.Error())
+			}
 
-	if wtbuf.String() != string(rdbuf) {
-		t.Fatalf("Read produced %q but WriteTo produced %q", string(rdbuf), wtbuf.String())
+			wtbuf := &bytes.Buffer{}
+			if _, err := r2.WriteTo(wtbuf); err != nil {
+				t.Fatalf("response.WriteTo(): %s", err.Error())
+			}
+
+			if wtbuf.String() != string(rdbuf) {
+				t.Fatalf("Read produced %q but WriteTo produced %q", string(rdbuf), wtbuf.String())
+			}
+		})
 	}
 }