Ensure the client is closed when closing the connection. (#47)

Author: ioquatix

lib/async/websocket/client.rb

@@ -12,6 +12,8 @@
 
 require 'async/http/client'
 
+require 'delegate'
+
 module Async
 	module WebSocket
 		# This is a basic synchronous websocket client:
@@ -31,13 +33,29 @@ def self.open(endpoint, **options, &block)
 				end
 			end
 
+			class ClientCloseDecorator < SimpleDelegator
+				def initialize(client, connection)
+					@client = client
+					super(connection)
+				end
+				
+				def close
+					super
+					
+					if @client
+						@client.close
+						@client = nil
+					end
+				end
+			end
+			
 			# @return [Connection] an open websocket connection to the given endpoint.
 			def self.connect(endpoint, *arguments, **options, &block)
 				client = self.open(endpoint, *arguments)
 				connection = client.connect(endpoint.authority, endpoint.path, **options)
-					
-				return connection unless block_given?
-					
+				
+				return ClientCloseDecorator.new(client, connection) unless block_given?
+				
 				begin
 					yield connection
 				ensure

test/async/websocket/client.rb

@@ -0,0 +1,37 @@
+
+
+ClientExamples = Sus::Shared("a websocket client") do
+	include Sus::Fixtures::Async::HTTP::ServerContext
+	
+	let(:app) do
+		Protocol::HTTP::Middleware.for do |request|
+			Async::WebSocket::Adapters::HTTP.open(request) do |connection|
+				while message = connection.read
+					connection.write(message)
+				end
+				
+				connection.close
+			end or Protocol::HTTP::Response[404, {}, []]
+		end
+	end
+	
+	let(:timeout) {nil}
+	
+	it "can connect to a websocket server and close underlying client" do
+		connection = Async::WebSocket::Client.connect(client_endpoint)
+		connection.send_text("Hello World!")
+		message = connection.read
+		expect(message.to_str).to be == "Hello World!"
+		connection.close
+	end
+end
+
+describe Async::WebSocket::Client do
+	with "h1", protocol: Async::HTTP::Protocol::HTTP1 do
+		it_behaves_like ClientExamples
+	end
+
+	with "h2", protocol: Async::HTTP::Protocol::HTTP2 do
+		it_behaves_like ClientExamples
+	end
+end

test/async/websocket/server.rb

@@ -11,7 +11,7 @@
 
 require 'sus/fixtures/async/http/server_context'
 
-WebSocketServerExamples = Sus::Shared('a websocket server') do
+ServerExamples = Sus::Shared('a websocket server') do
 	include Sus::Fixtures::Async::HTTP::ServerContext
 
 	let(:message) {"Hello World"}
@@ -72,11 +72,11 @@
 describe Async::HTTP::Protocol::HTTP1 do
 	let(:protocol) {subject}
 
-	it_behaves_like WebSocketServerExamples
+	it_behaves_like ServerExamples
 end
 
 describe Async::HTTP::Protocol::HTTP2 do
 	let(:protocol) {subject}
 
-	it_behaves_like WebSocketServerExamples
+	it_behaves_like ServerExamples
 end