github.com/simonmittag/ws@v1.1.0-rc.5.0.20210419231947-82b846128245/wsutil/handler_test.go (about)

     1  package wsutil
     2  
     3  import (
     4  	"bytes"
     5  	"runtime"
     6  	"testing"
     7  
     8  	"github.com/simonmittag/ws"
     9  )
    10  
    11  func TestControlHandler(t *testing.T) {
    12  	for _, test := range []struct {
    13  		name  string
    14  		state ws.State
    15  		in    ws.Frame
    16  		out   ws.Frame
    17  		noOut bool
    18  		err   error
    19  	}{
    20  		{
    21  			name: "ping",
    22  			in:   ws.NewPingFrame(nil),
    23  			out:  ws.NewPongFrame(nil),
    24  		},
    25  		{
    26  			name: "ping",
    27  			in:   ws.NewPingFrame([]byte("catch the ball")),
    28  			out:  ws.NewPongFrame([]byte("catch the ball")),
    29  		},
    30  		{
    31  			name:  "ping",
    32  			state: ws.StateServerSide,
    33  			in:    ws.MaskFrame(ws.NewPingFrame([]byte("catch the ball"))),
    34  			out:   ws.NewPongFrame([]byte("catch the ball")),
    35  		},
    36  		{
    37  			name:  "pong",
    38  			in:    ws.NewPongFrame(nil),
    39  			noOut: true,
    40  		},
    41  		{
    42  			name:  "pong",
    43  			in:    ws.NewPongFrame([]byte("catched")),
    44  			noOut: true,
    45  		},
    46  		{
    47  			name: "close",
    48  			in:   ws.NewCloseFrame(nil),
    49  			out:  ws.NewCloseFrame(nil),
    50  			err: ClosedError{
    51  				Code: ws.StatusNoStatusRcvd,
    52  			},
    53  		},
    54  		{
    55  			name: "close",
    56  			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
    57  				ws.StatusGoingAway, "goodbye!",
    58  			)),
    59  			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
    60  				ws.StatusGoingAway, "",
    61  			)),
    62  			err: ClosedError{
    63  				Code:   ws.StatusGoingAway,
    64  				Reason: "goodbye!",
    65  			},
    66  		},
    67  		{
    68  			name:  "close",
    69  			state: ws.StateServerSide,
    70  			in: ws.MaskFrame(ws.NewCloseFrame(ws.NewCloseFrameBody(
    71  				ws.StatusGoingAway, "goodbye!",
    72  			))),
    73  			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
    74  				ws.StatusGoingAway, "",
    75  			)),
    76  			err: ClosedError{
    77  				Code:   ws.StatusGoingAway,
    78  				Reason: "goodbye!",
    79  			},
    80  		},
    81  		{
    82  			name: "close",
    83  			in: ws.NewCloseFrame(ws.NewCloseFrameBody(
    84  				ws.StatusNormalClosure, string([]byte{0, 200}),
    85  			)),
    86  			out: ws.NewCloseFrame(ws.NewCloseFrameBody(
    87  				ws.StatusProtocolError, ws.ErrProtocolInvalidUTF8.Error(),
    88  			)),
    89  			err: ws.ErrProtocolInvalidUTF8,
    90  		},
    91  	} {
    92  		t.Run(test.name, func(t *testing.T) {
    93  			defer func() {
    94  				if err := recover(); err != nil {
    95  					stack := make([]byte, 4096)
    96  					n := runtime.Stack(stack, true)
    97  					t.Fatalf(
    98  						"panic recovered: %v\n%s",
    99  						err, stack[:n],
   100  					)
   101  				}
   102  			}()
   103  			var (
   104  				out = bytes.NewBuffer(nil)
   105  				in  = bytes.NewReader(test.in.Payload)
   106  			)
   107  			c := ControlHandler{
   108  				Src:   in,
   109  				Dst:   out,
   110  				State: test.state,
   111  			}
   112  
   113  			err := c.Handle(test.in.Header)
   114  			if err != test.err {
   115  				t.Errorf("unexpected error: %v; want %v", err, test.err)
   116  			}
   117  
   118  			if in.Len() != 0 {
   119  				t.Errorf("handler did not drained the input")
   120  			}
   121  
   122  			act := out.Bytes()
   123  			switch {
   124  			case len(act) == 0 && test.noOut:
   125  				return
   126  			case len(act) == 0 && !test.noOut:
   127  				t.Errorf("unexpected silence")
   128  			case len(act) > 0 && test.noOut:
   129  				t.Errorf("unexpected sent frame")
   130  			default:
   131  				exp := ws.MustCompileFrame(test.out)
   132  				if !bytes.Equal(act, exp) {
   133  					fa := ws.MustReadFrame(bytes.NewReader(act))
   134  					fe := ws.MustReadFrame(bytes.NewReader(exp))
   135  					t.Errorf(
   136  						"unexpected sent frame:\n\tact: %+v\n\texp: %+v\nbytes:\n\tact: %v\n\texp: %v",
   137  						fa, fe, act, exp,
   138  					)
   139  				}
   140  			}
   141  		})
   142  	}
   143  }