response: Implement Write method

2 files changed, 52 insertions, 9 deletions

diff --git a/response.go b/response.go
index b02b507..17e2155 100644
--- a/response.go
+++ b/response.go
@@ -3,6 +3,7 @@ package gemini
 import (
 	"bufio"
 	"crypto/tls"
+	"fmt"
 	"io"
 	"strconv"
 )
@@ -135,6 +136,24 @@ func (b *readCloserBody) Read(p []byte) (n int, err error) {
 	return b.ReadCloser.Read(p)
 }
 
+// Write writes r to w in the Gemini response format, including the
+// header and body.
+//
+// This method consults the Status, Meta, and Body fields of the response.
+// The Response Body is closed after it is sent.
+func (r *Response) Write(w io.Writer) error {
+	if _, err := fmt.Fprintf(w, "%02d %s\r\n", r.Status, r.Meta); err != nil {
+		return err
+	}
+	if r.Body != nil {
+		defer r.Body.Close()
+		if _, err := io.Copy(w, r.Body); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // A ResponseWriter interface is used by a Gemini handler to construct
 // a Gemini response.
 //
diff --git a/response_test.go b/response_test.go
index ab2484f..889e1ab 100644
--- a/response_test.go
+++ b/response_test.go
@@ -6,13 +6,14 @@ import (
 	"testing"
 )
 
-func TestReadResponse(t *testing.T) {
+func TestReadWriteResponse(t *testing.T) {
 	tests := []struct {
-		Raw    string
-		Status int
-		Meta   string
-		Body   string
-		Err    error
+		Raw       string
+		Status    int
+		Meta      string
+		Body      string
+		Err       error
+		SkipWrite bool
 	}{
 		{
 			Raw:    "20 text/gemini\r\nHello, world!\nWelcome to my capsule.",
@@ -31,9 +32,10 @@ func TestReadResponse(t *testing.T) {
 			Meta:   "/redirect",
 		},
 		{
-			Raw:    "31 /redirect\r\nThis body is ignored.",
-			Status: 31,
-			Meta:   "/redirect",
+			Raw:       "31 /redirect\r\nThis body is ignored.",
+			Status:    31,
+			Meta:      "/redirect",
+			SkipWrite: true, // skip write test since result won't match Raw
 		},
 		{
 			Raw:    "99 Unknown status code\r\n",
@@ -100,4 +102,26 @@ func TestReadResponse(t *testing.T) {
 			t.Errorf("expected body = %#v, got %#v", test.Body, body)
 		}
 	}
+
+	for _, test := range tests {
+		if test.Err != nil || test.SkipWrite {
+			continue
+		}
+		resp := &Response{
+			Status: test.Status,
+			Meta:   test.Meta,
+			Body:   io.NopCloser(strings.NewReader(test.Body)),
+		}
+
+		var b strings.Builder
+		if err := resp.Write(&b); err != nil {
+			t.Error(err)
+			continue
+		}
+
+		got := b.String()
+		if got != test.Raw {
+			t.Errorf("expected %#v, got %#v", test.Raw, got)
+		}
+	}
 }