add websocket_p method to Mojo::UserAgent

mojolicious · Dec 2, 2017 · 1f6f00c · 1f6f00c
1 parent b83ce9e
commit 1f6f00c
Show file tree

Hide file tree

Showing 3 changed files with 68 additions and 16 deletions.
diff --git a/Changes b/Changes
@@ -1,5 +1,6 @@
 
-7.58  2017-11-24
+7.58  2017-12-02
+  - Added websocket_p method to Mojo::UserAgent.
 
 7.57  2017-11-18
   - Fixed installation problems with some versions of Perl on Windows.

diff --git a/lib/Mojo/UserAgent.pm b/lib/Mojo/UserAgent.pm
@@ -66,27 +66,18 @@ sub start {
   return $tx;
 }
 
-sub start_p {
-  my ($self, $tx) = @_;
-
-  my $promise = Mojo::Promise->new;
-  $self->start(
-    $tx => sub {
-      my ($self, $tx) = @_;
-      my $err = $tx->error;
-      $promise->resolve($tx) if !$err || $err->{code};
-      $promise->reject($err->{message});
-    }
-  );
-
-  return $promise;
-}
+sub start_p { shift->_promise(0, @_) }
 
 sub websocket {
   my ($self, $cb) = (shift, pop);
   $self->start($self->build_websocket_tx(@_), $cb);
 }
 
+sub websocket_p {
+  my $self = shift;
+  return $self->_promise(1, $self->build_websocket_tx(@_));
+}
+
 sub _cleanup {
   my $self = shift;
   delete $self->{pid};
@@ -272,6 +263,24 @@ sub _finish {
   $c->{cb}($self, $old) unless $self->_redirect($c, $old);
 }
 
+sub _promise {
+  my ($self, $ws, $tx) = @_;
+
+  my $promise = Mojo::Promise->new;
+  $self->start(
+    $tx => sub {
+      my ($self, $tx) = @_;
+      $promise->reject('WebSocket handshake failed')
+        if $ws && !$tx->is_websocket;
+      my $err = $tx->error;
+      $promise->resolve($tx) if !$err || $err->{code};
+      $promise->reject($err->{message});
+    }
+  );
+
+  return $promise;
+}
+
 sub _read {
   my ($self, $id, $chunk) = @_;
 
@@ -1054,6 +1063,29 @@ but also increases memory usage by up to 300KiB per connection.
     'Sec-WebSocket-Extensions' => 'permessage-deflate'
   } => sub {...});
 
+=head2 websocket_p
+
+  my $promise = $ua->websocket_p('ws://example.com');
+
+Same as L</"websocket">, but returns a L<Mojo::Promise> object instead of
+accepting a callback.
+
+  $ua->websocket_p('wss://example.com/echo')->then(sub {
+    my $tx = shift;
+    my $promise = Mojo::Promise->new;
+    $tx->on(finish => sub { $promise->resolve });
+    $tx->on(message => sub {
+      my ($tx, $msg) = @_;
+      say "WebSocket message: $msg";
+      $tx->finish;
+    });
+    $tx->send('Hi!');
+    return $promise;
+  })->catch(sub {
+    my $err = shift;
+    warn "WebSocket error: $err";
+  })->wait;
+
 =head1 DEBUGGING
 
 You can set the C<MOJO_USERAGENT_DEBUG> environment variable to get some

diff --git a/t/mojo/websocket.t b/t/mojo/websocket.t
@@ -5,6 +5,7 @@ BEGIN { $ENV{MOJO_REACTOR} = 'Mojo::Reactor::Poll' }
 use Test::More;
 use Mojo::ByteStream 'b';
 use Mojo::IOLoop;
+use Mojo::Promise;
 use Mojo::Transaction::WebSocket;
 use Mojo::UserAgent;
 use Mojolicious::Lite;
@@ -333,6 +334,24 @@ $ua->websocket(
 Mojo::IOLoop->start;
 is $result, 'foo bar', 'right result';
 
+# Promises
+$result = undef;
+$ua->websocket_p('/trim')->then(
+  sub {
+    my $tx      = shift;
+    my $promise = Mojo::Promise->new;
+    $tx->on(finish => sub { $promise->resolve });
+    $tx->on(message => sub { shift->finish; $result = pop });
+    $tx->send(' also works! ');
+    return $promise;
+  }
+)->wait;
+is $result, 'also works!', 'right result';
+$result = undef;
+$ua->websocket_p('/foo')->then(sub { $result = 'test failed' })
+  ->catch(sub                      { $result = shift })->wait;
+is $result, 'WebSocket handshake failed', 'right result';
+
 # Dies
 ($ws, $code, $msg) = ();
 my $finished;