Repository: nhooyr/websocket
Pull request #159 (2019-10-05 23:35:57)

Add websocket.CloseStatus
Will make it much more convenient and less error prone to access
the StatusCode of a *CloseError.

Related #128 

Precedent for this is in https://godoc.org/gocloud.dev/gcerrors#Code

It's annoying to do the whole xerrors.As dance when you only want to access the status code.

I was only blocked on the name for this but I think `CloseStatus` is nice.

@ammar @coadler 

<!--
Please read the contributing guidelines.
https://github.com/nhooyr/websocket/blob/master/.github/CONTRIBUTING.md#pull-requests
-->



Patch:
diff --git a/README.md b/README.md
index 3ab395e..f6afbd8 100644
--- a/README.md
+++ b/README.md
@@ -33,7 +33,9 @@ go get nhooyr.io/websocket
 
 For a production quality example that shows off the full API, see the [echo example on the godoc](https://godoc.org/nhooyr.io/websocket#example-package--Echo). On github, the example is at [example_echo_test.go](./example_echo_test.go).
 
-Use the [errors.As](https://golang.org/pkg/errors/#As) function [new in Go 1.13](https://golang.org/doc/go1.13#error_wrapping) to check for [websocket.CloseError](https://godoc.org/nhooyr.io/websocket#CloseError). See the [CloseError godoc example](https://godoc.org/nhooyr.io/websocket#example-CloseError).
+Use the [errors.As](https://golang.org/pkg/errors/#As) function [new in Go 1.13](https://golang.org/doc/go1.13#error_wrapping) to check for [websocket.CloseError](https://godoc.org/nhooyr.io/websocket#CloseError).
+There is also [websocket.CloseStatus](https://godoc.org/nhooyr.io/websocket#CloseStatus) to quickly grab the close status code out of a [websocket.CloseError](https://godoc.org/nhooyr.io/websocket#CloseError).
+See the [CloseError godoc example](https://godoc.org/nhooyr.io/websocket#example-CloseError).
 
 ### Server
 
diff --git a/doc.go b/doc.go
index a17bfb0..1610eed 100644
--- a/doc.go
+++ b/doc.go
@@ -16,6 +16,7 @@
 // comparison with existing implementations.
 //
 // Use the errors.As function new in Go 1.13 to check for websocket.CloseError.
+// Or use the CloseStatus function to grab the StatusCode out of a websocket.CloseError
 // See the CloseError example.
 //
 // Wasm
diff --git a/frame.go b/frame.go
index 4b170c5..b5301d7 100644
--- a/frame.go
+++ b/frame.go
@@ -2,6 +2,7 @@ package websocket
 
 import (
 	"encoding/binary"
+	"errors"
 	"fmt"
 	"io"
 	"math"
@@ -252,6 +253,17 @@ func (ce CloseError) Error() string {
 	return fmt.Sprintf("status = %v and reason = %q", ce.Code, ce.Reason)
 }
 
+// CloseStatus is a convenience wrapper around errors.As to grab
+// the status code from a *CloseError. If the passed error is nil
+// or not a *CloseError, the returned StatusCode will be -1.
+func CloseStatus(err error) StatusCode {
+	var ce CloseError
+	if errors.As(err, &ce) {
+		return ce.Code
+	}
+	return -1
+}
+
 func parseClosePayload(p []byte) (CloseError, error) {
 	if len(p) == 0 {
 		return CloseError{
diff --git a/internal/assert/assert.go b/internal/assert/assert.go
new file mode 100644
index 0000000..e57abfd
--- /dev/null
+++ b/internal/assert/assert.go
@@ -0,0 +1,63 @@
+package assert
+
+import (
+	"fmt"
+	"reflect"
+
+	"github.com/google/go-cmp/cmp"
+)
+
+// https://github.com/google/go-cmp/issues/40#issuecomment-328615283
+func cmpDiff(exp, act interface{}) string {
+	return cmp.Diff(exp, act, deepAllowUnexported(exp, act))
+}
+
+func deepAllowUnexported(vs ...interface{}) cmp.Option {
+	m := make(map[reflect.Type]struct{})
+	for _, v := range vs {
+		structTypes(reflect.ValueOf(v), m)
+	}
+	var typs []interface{}
+	for t := range m {
+		typs = append(typs, reflect.New(t).Elem().Interface())
+	}
+	return cmp.AllowUnexported(typs...)
+}
+
+func structTypes(v reflect.Value, m map[reflect.Type]struct{}) {
+	if !v.IsValid() {
+		return
+	}
+	switch v.Kind() {
+	case reflect.Ptr:
+		if !v.IsNil() {
+			structTypes(v.Elem(), m)
+		}
+	case reflect.Interface:
+		if !v.IsNil() {
+			structTypes(v.Elem(), m)
+		}
+	case reflect.Slice, reflect.Array:
+		for i := 0; i < v.Len(); i++ {
+			structTypes(v.Index(i), m)
+		}
+	case reflect.Map:
+		for _, k := range v.MapKeys() {
+			structTypes(v.MapIndex(k), m)
+		}
+	case reflect.Struct:
+		m[v.Type()] = struct{}{}
+		for i := 0; i < v.NumField(); i++ {
+			structTypes(v.Field(i), m)
+		}
+	}
+}
+
+// Equalf compares exp to act and if they are not equal, returns
+// an error describing an error.
+func Equalf(exp, act interface{}, f string, v ...interface{}) error {
+	if diff := cmpDiff(exp, act); diff != "" {
+		return fmt.Errorf(f+": %v", append(v, diff)...)
+	}
+	return nil
+}


Test patch:
diff --git a/assert_test.go b/assert_test.go
index e67ed53..26fd1d4 100644
--- a/assert_test.go
+++ b/assert_test.go
@@ -2,15 +2,12 @@ package websocket_test
 
 import (
 	"context"
-	"fmt"
 	"math/rand"
-	"reflect"
 	"strings"
 	"time"
 
-	"github.com/google/go-cmp/cmp"
-
 	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/internal/assert"
 	"nhooyr.io/websocket/wsjson"
 )
 
@@ -18,59 +15,6 @@ func init() {
 	rand.Seed(time.Now().UnixNano())
 }
 
-// https://github.com/google/go-cmp/issues/40#issuecomment-328615283
-func cmpDiff(exp, act interface{}) string {
-	return cmp.Diff(exp, act, deepAllowUnexported(exp, act))
-}
-
-func deepAllowUnexported(vs ...interface{}) cmp.Option {
-	m := make(map[reflect.Type]struct{})
-	for _, v := range vs {
-		structTypes(reflect.ValueOf(v), m)
-	}
-	var typs []interface{}
-	for t := range m {
-		typs = append(typs, reflect.New(t).Elem().Interface())
-	}
-	return cmp.AllowUnexported(typs...)
-}
-
-func structTypes(v reflect.Value, m map[reflect.Type]struct{}) {
-	if !v.IsValid() {
-		return
-	}
-	switch v.Kind() {
-	case reflect.Ptr:
-		if !v.IsNil() {
-			structTypes(v.Elem(), m)
-		}
-	case reflect.Interface:
-		if !v.IsNil() {
-			structTypes(v.Elem(), m)
-		}
-	case reflect.Slice, reflect.Array:
-		for i := 0; i < v.Len(); i++ {
-			structTypes(v.Index(i), m)
-		}
-	case reflect.Map:
-		for _, k := range v.MapKeys() {
-			structTypes(v.MapIndex(k), m)
-		}
-	case reflect.Struct:
-		m[v.Type()] = struct{}{}
-		for i := 0; i < v.NumField(); i++ {
-			structTypes(v.Field(i), m)
-		}
-	}
-}
-
-func assertEqualf(exp, act interface{}, f string, v ...interface{}) error {
-	if diff := cmpDiff(exp, act); diff != "" {
-		return fmt.Errorf(f+": %v", append(v, diff)...)
-	}
-	return nil
-}
-
 func assertJSONEcho(ctx context.Context, c *websocket.Conn, n int) error {
 	exp := randString(n)
 	err := wsjson.Write(ctx, c, exp)
@@ -84,7 +28,7 @@ func assertJSONEcho(ctx context.Context, c *websocket.Conn, n int) error {
 		return err
 	}
 
-	return assertEqualf(exp, act, "unexpected JSON")
+	return assert.Equalf(exp, act, "unexpected JSON")
 }
 
 func assertJSONRead(ctx context.Context, c *websocket.Conn, exp interface{}) error {
@@ -94,7 +38,7 @@ func assertJSONRead(ctx context.Context, c *websocket.Conn, exp interface{}) err
 		return err
 	}
 
-	return assertEqualf(exp, act, "unexpected JSON")
+	return assert.Equalf(exp, act, "unexpected JSON")
 }
 
 func randBytes(n int) []byte {
@@ -126,13 +70,13 @@ func assertEcho(ctx context.Context, c *websocket.Conn, typ websocket.MessageTyp
 	if err != nil {
 		return err
 	}
-	err = assertEqualf(typ, typ2, "unexpected data type")
+	err = assert.Equalf(typ, typ2, "unexpected data type")
 	if err != nil {
 		return err
 	}
-	return assertEqualf(p, p2, "unexpected payload")
+	return assert.Equalf(p, p2, "unexpected payload")
 }
 
 func assertSubprotocol(c *websocket.Conn, exp string) error {
-	return assertEqualf(exp, c.Subprotocol(), "unexpected subprotocol")
+	return assert.Equalf(exp, c.Subprotocol(), "unexpected subprotocol")
 }
diff --git a/conn_test.go b/conn_test.go
index 12788c3..0e012bf 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -32,6 +32,7 @@ import (
 	"go.uber.org/multierr"
 
 	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/internal/assert"
 	"nhooyr.io/websocket/internal/wsecho"
 	"nhooyr.io/websocket/wsjson"
 	"nhooyr.io/websocket/wspb"
@@ -127,7 +128,7 @@ func TestHandshake(t *testing.T) {
 				if err != nil {
 					return fmt.Errorf("request is missing mycookie: %w", err)
 				}
-				err = assertEqualf("myvalue", cookie.Value, "unexpected cookie value")
+				err = assert.Equalf("myvalue", cookie.Value, "unexpected cookie value")
 				if err != nil {
 					return err
 				}
@@ -219,7 +220,7 @@ func TestConn(t *testing.T) {
 				}
 				for h, exp := range headers {
 					value := resp.Header.Get(h)
-					err := assertEqualf(exp, value, "unexpected value for header %v", h)
+					err := assert.Equalf(exp, value, "unexpected value for header %v", h)
 					if err != nil {
 						return err
 					}
@@ -276,11 +277,11 @@ func TestConn(t *testing.T) {
 				time.Sleep(1)
 				nc.SetWriteDeadline(time.Now().Add(time.Second * 15))
 
-				err := assertEqualf(websocket.Addr{}, nc.LocalAddr(), "net conn local address is not equal to websocket.Addr")
+				err := assert.Equalf(websocket.Addr{}, nc.LocalAddr(), "net conn local address is not equal to websocket.Addr")
 				if err != nil {
 					return err
 				}
-				err = assertEqualf(websocket.Addr{}, nc.RemoteAddr(), "net conn remote address is not equal to websocket.Addr")
+				err = assert.Equalf(websocket.Addr{}, nc.RemoteAddr(), "net conn remote address is not equal to websocket.Addr")
 				if err != nil {
 					return err
 				}
@@ -310,13 +311,13 @@ func TestConn(t *testing.T) {
 
 				// Ensure the close frame is converted to an EOF and multiple read's after all return EOF.
 				err2 := assertNetConnRead(nc, "hello")
-				err := assertEqualf(io.EOF, err2, "unexpected error")
+				err := assert.Equalf(io.EOF, err2, "unexpected error")
 				if err != nil {
 					return err
 				}
 
 				err2 = assertNetConnRead(nc, "hello")
-				return assertEqualf(io.EOF, err2, "unexpected error")
+				return assert.Equalf(io.EOF, err2, "unexpected error")
 			},
 		},
 		{
@@ -585,8 +586,8 @@ func TestConn(t *testing.T) {
 					return err
 				}
 				_, _, err = c.Read(ctx)
-				cerr := &websocket.CloseError{}
-				if !errors.As(err, cerr) || cerr.Code != websocket.StatusProtocolError {
+				var cerr websocket.CloseError
+				if !errors.As(err, &cerr) || cerr.Code != websocket.StatusProtocolError {
 					return fmt.Errorf("expected close error with StatusProtocolError: %+v", err)
 				}
 				return nil
@@ -772,7 +773,7 @@ func TestConn(t *testing.T) {
 				if err != nil {
 					return err
 				}
-				err = assertEqualf("hi", v, "unexpected JSON")
+				err = assert.Equalf("hi", v, "unexpected JSON")
 				if err != nil {
 					return err
 				}
@@ -780,7 +781,7 @@ func TestConn(t *testing.T) {
 				if err != nil {
 					return err
 				}
-				return assertEqualf("hi", string(b), "unexpected JSON")
+				return assert.Equalf("hi", string(b), "unexpected JSON")
 			},
 			client: func(ctx context.Context, c *websocket.Conn) error {
 				err := wsjson.Write(ctx, c, "hi")
@@ -1079,11 +1080,11 @@ func TestAutobahn(t *testing.T) {
 					if err != nil {
 						return err
 					}
-					err = assertEqualf(typ, actTyp, "unexpected message type")
+					err = assert.Equalf(typ, actTyp, "unexpected message type")
 					if err != nil {
 						return err
 					}
-					return assertEqualf(p, p2, "unexpected message")
+					return assert.Equalf(p, p2, "unexpected message")
 				})
 			}
 		}
@@ -1859,7 +1860,7 @@ func assertCloseStatus(err error, code websocket.StatusCode) error {
 	if !errors.As(err, &cerr) {
 		return fmt.Errorf("no websocket close error in error chain: %+v", err)
 	}
-	return assertEqualf(code, cerr.Code, "unexpected status code")
+	return assert.Equalf(code, cerr.Code, "unexpected status code")
 }
 
 func assertProtobufRead(ctx context.Context, c *websocket.Conn, exp interface{}) error {
@@ -1871,7 +1872,7 @@ func assertProtobufRead(ctx context.Context, c *websocket.Conn, exp interface{})
 		return err
 	}
 
-	return assertEqualf(exp, act, "unexpected protobuf")
+	return assert.Equalf(exp, act, "unexpected protobuf")
 }
 
 func assertNetConnRead(r io.Reader, exp string) error {
@@ -1880,7 +1881,7 @@ func assertNetConnRead(r io.Reader, exp string) error {
 	if err != nil {
 		return err
 	}
-	return assertEqualf(exp, string(act), "unexpected net conn read")
+	return assert.Equalf(exp, string(act), "unexpected net conn read")
 }
 
 func assertErrorContains(err error, exp string) error {
@@ -1902,11 +1903,11 @@ func assertReadFrame(ctx context.Context, c *websocket.Conn, opcode websocket.Op
 	if err != nil {
 		return err
 	}
-	err = assertEqualf(opcode, actOpcode, "unexpected frame opcode with payload %q", actP)
+	err = assert.Equalf(opcode, actOpcode, "unexpected frame opcode with payload %q", actP)
 	if err != nil {
 		return err
 	}
-	return assertEqualf(p, actP, "unexpected frame %v payload", opcode)
+	return assert.Equalf(p, actP, "unexpected frame %v payload", opcode)
 }
 
 func assertReadCloseFrame(ctx context.Context, c *websocket.Conn, code websocket.StatusCode) error {
@@ -1914,7 +1915,7 @@ func assertReadCloseFrame(ctx context.Context, c *websocket.Conn, code websocket
 	if err != nil {
 		return err
 	}
-	err = assertEqualf(websocket.OpClose, actOpcode, "unexpected frame opcode with payload %q", actP)
+	err = assert.Equalf(websocket.OpClose, actOpcode, "unexpected frame opcode with payload %q", actP)
 	if err != nil {
 		return err
 	}
@@ -1922,7 +1923,7 @@ func assertReadCloseFrame(ctx context.Context, c *websocket.Conn, code websocket
 	if err != nil {
 		return fmt.Errorf("failed to parse close frame payload: %w", err)
 	}
-	return assertEqualf(ce.Code, code, "unexpected frame close frame code with payload %q", actP)
+	return assert.Equalf(ce.Code, code, "unexpected frame close frame code with payload %q", actP)
 }
 
 func assertCloseHandshake(ctx context.Context, c *websocket.Conn, code websocket.StatusCode, reason string) error {
@@ -1960,11 +1961,11 @@ func assertReadMessage(ctx context.Context, c *websocket.Conn, typ websocket.Mes
 	if err != nil {
 		return err
 	}
-	err = assertEqualf(websocket.MessageText, actTyp, "unexpected frame opcode with payload %q", actP)
+	err = assert.Equalf(websocket.MessageText, actTyp, "unexpected frame opcode with payload %q", actP)
 	if err != nil {
 		return err
 	}
-	return assertEqualf(p, actP, "unexpected frame %v payload", actTyp)
+	return assert.Equalf(p, actP, "unexpected frame %v payload", actTyp)
 }
 
 func BenchmarkConn(b *testing.B) {
diff --git a/example_test.go b/example_test.go
index 2cedddf..1cb3d79 100644
--- a/example_test.go
+++ b/example_test.go
@@ -4,7 +4,6 @@ package websocket_test
 
 import (
 	"context"
-	"errors"
 	"log"
 	"net/http"
 	"time"
@@ -76,8 +75,7 @@ func ExampleCloseError() {
 	defer c.Close(websocket.StatusInternalError, "the sky is falling")
 
 	_, _, err = c.Reader(ctx)
-	var cerr websocket.CloseError
-	if !errors.As(err, &cerr) || cerr.Code != websocket.StatusNormalClosure {
+	if websocket.CloseStatus(err) != websocket.StatusNormalClosure {
 		log.Fatalf("expected to be disconnected with StatusNormalClosure but got: %+v", err)
 		return
 	}
diff --git a/frame_test.go b/frame_test.go
index 7d2a571..8c5760e 100644
--- a/frame_test.go
+++ b/frame_test.go
@@ -13,6 +13,8 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+
+	"nhooyr.io/websocket/internal/assert"
 )
 
 func init() {
@@ -376,3 +378,43 @@ func BenchmarkXOR(b *testing.B) {
 		})
 	}
 }
+
+func TestCloseStatus(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		name string
+		in   error
+		exp  StatusCode
+	}{
+		{
+			name: "nil",
+			in:   nil,
+			exp:  -1,
+		},
+		{
+			name: "io.EOF",
+			in:   io.EOF,
+			exp:  -1,
+		},
+		{
+			name: "StatusInternalError",
+			in: CloseError{
+				Code: StatusInternalError,
+			},
+			exp: StatusInternalError,
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			err := assert.Equalf(tc.exp, CloseStatus(tc.in), "unexpected close status")
+			if err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+}
diff --git a/websocket_js_test.go b/websocket_js_test.go
index a3bb763..9808e70 100644
--- a/websocket_js_test.go
+++ b/websocket_js_test.go
@@ -8,6 +8,7 @@ import (
 	"time"
 
 	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/internal/assert"
 )
 
 func TestConn(t *testing.T) {
@@ -29,7 +30,7 @@ func TestConn(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	err = assertEqualf(&http.Response{}, resp, "unexpected http response")
+	err = assert.Equalf(&http.Response{}, resp, "unexpected http response")
 	if err != nil {
 		t.Fatal(err)
 	}