improved request_ok method in Test::Mojo to handle WebSocket handshakes

mojolicious · Jan 18, 2014 · 4154f72 · 4154f72
1 parent ef84ef1
commit 4154f72
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 41 deletions.
diff --git a/Changes b/Changes
@@ -8,6 +8,7 @@
   - Removed deprecated app_url method from Mojo::UserAgent.
   - Removed deprecated detect_proxy method from Mojo::UserAgent.
   - Removed deprecated need_proxy method from Mojo::UserAgent.
+  - Improved request_ok method in Test::Mojo to handle WebSocket handshakes.
   - Improved Mojo::IOLoop::Server to use address and port for descriptor
     inheritance.
   - Improved list of available commands to be alphabetical. (jberger)

diff --git a/lib/Test/Mojo.pm b/lib/Test/Mojo.pm
@@ -87,7 +87,7 @@ sub content_type_unlike {
     $regex, $desc);
 }
 
-sub delete_ok { shift->_request_ok(delete => @_) }
+sub delete_ok { shift->_build_ok(DELETE => @_) }
 
 sub element_exists {
   my ($self, $selector, $desc) = @_;
@@ -116,8 +116,8 @@ sub finished_ok {
   return $self->_test('ok', $ok, "WebSocket closed with status $code");
 }
 
-sub get_ok  { shift->_request_ok(get  => @_) }
-sub head_ok { shift->_request_ok(head => @_) }
+sub get_ok  { shift->_build_ok(GET  => @_) }
+sub head_ok { shift->_build_ok(HEAD => @_) }
 
 sub header_is {
   my ($self, $name, $value, $desc) = @_;
@@ -211,23 +211,19 @@ sub message_unlike {
   return $self->_message('unlike', $regex, $desc || 'message is not similar');
 }
 
-sub options_ok { shift->_request_ok(options => @_) }
+sub options_ok { shift->_build_ok(OPTIONS => @_) }
 
 sub or {
   my ($self, $cb) = @_;
   $self->$cb unless $self->success;
   return $self;
 }
 
-sub patch_ok { shift->_request_ok(patch => @_) }
-sub post_ok  { shift->_request_ok(post  => @_) }
-sub put_ok   { shift->_request_ok(put   => @_) }
+sub patch_ok { shift->_build_ok(PATCH => @_) }
+sub post_ok  { shift->_build_ok(POST  => @_) }
+sub put_ok   { shift->_build_ok(PUT   => @_) }
 
-sub request_ok {
-  my $self = shift;
-  my $tx   = $self->tx($self->ua->start(shift))->tx;
-  return $self->_test('ok', $tx->is_finished, shift || 'perform request');
-}
+sub request_ok { shift->_request_ok(@_) }
 
 sub reset_session {
   my $self = shift;
@@ -280,25 +276,14 @@ sub text_unlike {
 }
 
 sub websocket_ok {
-  my ($self, $url) = (shift, shift);
-
-  # Establish WebSocket connection
-  $self->{messages} = [];
-  $self->{finished} = undef;
-  $self->ua->websocket(
-    $url => @_ => sub {
-      my ($ua, $tx) = @_;
-      $self->tx($tx);
-      $tx->on(finish => sub { shift; $self->{finished} = [@_] });
-      $tx->on(binary => sub { push @{$self->{messages}}, [binary => pop] });
-      $tx->on(text   => sub { push @{$self->{messages}}, [text   => pop] });
-      Mojo::IOLoop->stop;
-    }
-  );
-  Mojo::IOLoop->start;
+  my $self = shift;
+  return $self->_request_ok($self->ua->build_websocket_tx(@_));
+}
 
-  my $desc = encode 'UTF-8', "WebSocket $url";
-  return $self->_test('ok', $self->tx->is_websocket, $desc);
+sub _build_ok {
+  my $self = shift;
+  local $Test::Builder::Level = $Test::Builder::Level + 1;
+  return $self->_request_ok($self->ua->build_tx(@_));
 }
 
 sub _json {
@@ -326,14 +311,36 @@ sub _message {
 }
 
 sub _request_ok {
-  my ($self, $method, $url) = (shift, shift, shift);
+  my ($self, $tx) = @_;
 
-  # Perform request against application
-  $self->tx($self->ua->$method($url, @_));
   local $Test::Builder::Level = $Test::Builder::Level + 1;
+
+  # Establish WebSocket connection
+  if (lc($tx->req->headers->upgrade // '') eq 'websocket') {
+    $self->{messages} = [];
+    $self->{finished} = undef;
+    $self->ua->start(
+      $tx => sub {
+        my ($ua, $tx) = @_;
+        $self->tx($tx);
+        $tx->on(finish => sub { shift; $self->{finished} = [@_] });
+        $tx->on(binary => sub { push @{$self->{messages}}, [binary => pop] });
+        $tx->on(text   => sub { push @{$self->{messages}}, [text   => pop] });
+        Mojo::IOLoop->stop;
+      }
+    );
+    Mojo::IOLoop->start;
+
+    my $desc = encode 'UTF-8', "WebSocket @{[$tx->req->url]}";
+    return $self->_test('ok', $self->tx->is_websocket, $desc);
+  }
+
+  # Perform request against application
+  $self->tx($self->ua->start($tx));
   my ($err, $code) = $self->tx->error;
   Test::More::diag $err if !(my $ok = !$err || $code) && $err;
-  return $self->_test('ok', $ok, encode('UTF-8', "@{[uc $method]} $url"));
+  return $self->_test('ok', $ok,
+    encode('UTF-8', "@{[uc $tx->req->method]} @{[$tx->req->url]}"));
 }
 
 sub _test {
@@ -821,6 +828,11 @@ Perform request and check for transport errors.
   my $tx = $t->ua->build_tx(FOO => '/test.json' => json => {foo => 1});
   $t->request_ok($tx)->status_is(200)->json_is({success => 1});
 
+  # WebSocket handshake without extensions
+  my $tx = $t->ua->build_websocket_tx('/foo');
+  $tx->req->headers->remove('Sec-WebSocket-Extensions');
+  $t->request_ok($tx)->message_ok->message_is('bar')->finish_ok;
+
 =head2 reset_session
 
   $t = $t->reset_session;

diff --git a/t/mojolicious/websocket_lite_app.t b/t/mojolicious/websocket_lite_app.t
@@ -151,16 +151,19 @@ $t->websocket_ok('/echo')->send_ok(0)->message_ok->message_is('echo: 0')
   ->finished_ok(1000);
 
 # 64bit binary message (extended limit and no compression)
-$t->websocket_ok('/echo' => {'Sec-WebSocket-Extensions' => 'nothing'});
+my $tx = $t->ua->build_websocket_tx('/echo');
+$tx->req->headers->remove('Sec-WebSocket-Extensions');
+$t->request_ok($tx);
 is $t->tx->max_websocket_size, 262144, 'right size';
 $t->tx->max_websocket_size(262145);
 $t->send_ok({binary => 'a' x 262145})
   ->message_ok->message_is({binary => 'a' x 262145})
   ->finish_ok->finished_ok(1005);
 
 # 64bit binary message (too large and no compression)
-$t->websocket_ok('/echo' => {'Sec-WebSocket-Extensions' => 'nothing'})
-  ->send_ok({binary => 'b' x 262145})->finished_ok(1009);
+$tx = $t->ua->build_websocket_tx('/echo');
+$tx->req->headers->remove('Sec-WebSocket-Extensions');
+$t->request_ok($tx)->send_ok({binary => 'b' x 262145})->finished_ok(1009);
 
 # Compressed message ("permessage-deflate")
 $t->websocket_ok('/echo');
@@ -249,7 +252,9 @@ $t->send_ok({binary => 'a' x 500})
 
 # Binary frame and events (no compression)
 my $bytes = b("I ♥ Mojolicious")->encode('UTF-16LE')->to_string;
-$t->websocket_ok('/bytes' => {'Sec-WebSocket-Extensions' => 'nothing'})
+$tx = $t->ua->build_websocket_tx('/bytes');
+$tx->req->headers->remove('Sec-WebSocket-Extensions');
+$t->request_ok($tx)
   ->header_isnt('Sec-WebSocket-Extensions' => 'permessage-deflate');
 ok !$t->tx->compressed, 'WebSocket has no compression';
 my $binary;
@@ -270,9 +275,11 @@ $t->finish_ok(1000 => 'Have a nice day!');
 is_deeply $close, [1000, 'Have a nice day!'], 'right status and message';
 
 # Binary roundtrips (no compression)
-$t->websocket_ok('/bytes' => {'Sec-WebSocket-Extensions' => 'nothing'})
-  ->send_ok({binary => $bytes})->message_ok->message_is($bytes)
-  ->send_ok({binary => $bytes})->message_ok->message_is($bytes)->finish_ok;
+$tx = $t->ua->build_websocket_tx('/bytes');
+$tx->req->headers->remove('Sec-WebSocket-Extensions');
+$t->request_ok($tx)->send_ok({binary => $bytes})
+  ->message_ok->message_is($bytes)->send_ok({binary => $bytes})
+  ->message_ok->message_is($bytes)->finish_ok;
 
 # Two responses
 $t->websocket_ok('/once')->send_ok('hello')