better WebSocket handshake tests

mojolicious · Jun 11, 2012 · e9df221 · e9df221
1 parent b5e9967
commit e9df221
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/lib/Mojo/Transaction/WebSocket.pm b/lib/Mojo/Transaction/WebSocket.pm
@@ -124,8 +124,7 @@ sub parse_frame {
   my ($self, $buffer) = @_;
 
   # Head
-  my $clone = $$buffer;
-  return unless length $clone >= 2;
+  return unless length(my $clone = $$buffer) >= 2;
   my $head = substr $clone, 0, 2;
 
   # FIN
@@ -149,8 +148,7 @@ sub parse_frame {
   elsif ($len == 126) {
     return unless length $clone > 4;
     $hlen = 4;
-    my $ext = substr $clone, 2, 2;
-    $len = unpack 'n', $ext;
+    $len = unpack 'n', substr($clone, 2, 2);
     warn "-- Extended 16bit payload ($len)\n" if DEBUG;
   }
 
@@ -222,14 +220,10 @@ sub server_handshake {
   my $self = shift;
 
   # WebSocket handshake
-  my $res         = $self->res;
-  my $res_headers = $res->headers;
-  $res->code(101);
-  $res_headers->upgrade('websocket');
-  $res_headers->connection('Upgrade');
+  my $res_headers = $self->res->code(101)->headers;
+  $res_headers->upgrade('websocket')->connection('Upgrade');
   my $req_headers = $self->req->headers;
-  my $protocol = $req_headers->sec_websocket_protocol || '';
-  $protocol =~ /^\s*([^,]+)/;
+  ($req_headers->sec_websocket_protocol || '') =~ /^\s*([^,]+)/;
   $res_headers->sec_websocket_protocol($1) if $1;
   $res_headers->sec_websocket_accept(
     $self->_challenge($req_headers->sec_websocket_key));

diff --git a/t/mojolicious/websocket_lite_app.t b/t/mojolicious/websocket_lite_app.t
@@ -10,7 +10,7 @@ BEGIN {
   $ENV{MOJO_REACTOR} = 'Mojo::Reactor::Poll';
 }
 
-use Test::More tests => 91;
+use Test::More tests => 96;
 
 use Mojo::ByteStream 'b';
 use Mojolicious::Lite;
@@ -115,6 +115,11 @@ $t->websocket_ok('/echo')->send_ok('hello again')
   ->message_is('echo: hello again')->send_ok('and one more time')
   ->message_is('echo: and one more time')->finish_ok;
 
+# WebSocket /echo (with custom protocol)
+$t->websocket_ok('/echo', {'Sec-WebSocket-Protocol' => 'foo, bar, baz'})
+  ->header_is('Sec-WebSocket-Protocol', 'foo')->send_ok('hello')
+  ->message_is('echo: hello')->finish_ok;
+
 # WebSocket /echo (zero)
 $t->websocket_ok('/echo')->send_ok(0)->message_is('echo: 0')->finish_ok;