From 7709299e541b240943191dfbca41abf528fda51a Mon Sep 17 00:00:00 2001
From: Samuel Williams <samuel.williams@oriontransfer.co.nz>
Date: Fri, 3 Feb 2023 14:35:39 +1300
Subject: [PATCH] Ensure the client is closed when closing the connection.

---
 lib/async/websocket/client.rb  | 24 +++++++++++++++++++---
 test/async/websocket/client.rb | 37 ++++++++++++++++++++++++++++++++++
 test/async/websocket/server.rb |  6 +++---
 3 files changed, 61 insertions(+), 6 deletions(-)
 create mode 100644 test/async/websocket/client.rb

diff --git a/lib/async/websocket/client.rb b/lib/async/websocket/client.rb
index 7d6d802..d9fb9eb 100644
--- a/lib/async/websocket/client.rb
+++ b/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
diff --git a/test/async/websocket/client.rb b/test/async/websocket/client.rb
new file mode 100644
index 0000000..2bd4ff2
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/test/async/websocket/server.rb b/test/async/websocket/server.rb
index 8ca19d8..5683971 100644
--- a/test/async/websocket/server.rb
+++ b/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