diff --git a/lib/sinatra/async.rb b/lib/sinatra/async.rb index 8ba5662..409f94e 100644 --- a/lib/sinatra/async.rb +++ b/lib/sinatra/async.rb @@ -136,6 +136,14 @@ def ahalt(*args) invoke { error_block! response.status } body response.body end + + # The given block will be executed if the user closes the connection + # prematurely (before we've sent a response). This is good for + # deregistering callbacks that would otherwise send the body (for + # example channel subscriptions). + def on_close(&blk) + env['async.close'].callback(&blk) + end end def self.registered(app) #:nodoc: diff --git a/lib/sinatra/async/test.rb b/lib/sinatra/async/test.rb index e20eca1..db7e7de 100644 --- a/lib/sinatra/async/test.rb +++ b/lib/sinatra/async/test.rb @@ -8,10 +8,31 @@ def async? end class Sinatra::Async::Test + class AsyncSession < Rack::MockSession + class AsyncCloser + def initialize + @callbacks, @errbacks = [], [] + end + def callback(&b) + @callbacks << b + end + def errback(&b) + @errbacks << b + end + def fail + @errbacks.each { |cb| cb.call } + @errbacks.clear + end + def succeed + @callbacks.each { |cb| cb.call } + @callbacks.clear + end + end + def request(uri, env) env['async.callback'] = lambda { |r| s,h,b = *r; handle_last_response(uri, env, s,h,b) } - env['async.close'] = lambda { raise 'close connection' } # XXX deal with this + env['async.close'] = AsyncCloser.new catch(:async) { super } @last_response ||= Rack::MockResponse.new(-1, {}, [], env["rack.errors"].flush) end @@ -48,6 +69,12 @@ def assert_async assert last_response.async? end + # Simulate a user closing the connection before a response is sent. + def async_close + raise ArgumentError, 'please make a request first' unless last_request + current_session.last_request.env['async.close'].succeed + end + def async_continue while b = app.options.async_schedules.shift b.call diff --git a/test/test_async.rb b/test/test_async.rb index 7e6b158..42d6124 100644 --- a/test/test_async.rb +++ b/test/test_async.rb @@ -12,6 +12,12 @@ class TestApp < Sinatra::Base set :environment, :test register Sinatra::Async + # Hack for storing some global data accessible in tests (normally you + # shouldn't need to do this!) + def self.singletons + @singletons ||= [] + end + error 401 do '401' end @@ -47,6 +53,18 @@ class TestApp < Sinatra::Base aget '/a401' do ahalt 401 end + + aget '/async_close' do + # don't call body here, the 'user' is going to 'disconnect' before we do + env['async.close'].callback { self.class.singletons << 'async_closed' } + end + + aget '/on_close' do + # sugared version of the above + on_close do + self.class.singletons << 'async_close_cleaned_up' + end + end end def app @@ -113,4 +131,20 @@ def test_error_blocks_async assert_equal 401, last_response.status assert_equal '401', last_response.body end + + def test_async_close + get '/async_close' + assert_async + async_continue + async_close + assert_equal 'async_closed', TestApp.singletons.shift + end + + def test_on_close + get '/on_close' + assert_async + async_continue + async_close + assert_equal 'async_close_cleaned_up', TestApp.singletons.shift + end end \ No newline at end of file