aboutsummaryrefslogtreecommitdiffstats
path: root/actionpack/test/controller/live_stream_test.rb
blob: 0c3884cd387a520e905c82b1508718ac050ddcdb (plain) (blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
require 'abstract_unit'
require 'concurrent/atomic/count_down_latch'
Thread.abort_on_exception = true

module ActionController
  class SSETest < ActionController::TestCase
    class SSETestController < ActionController::Base
      include ActionController::Live

      def basic_sse
        response.headers['Content-Type'] = 'text/event-stream'
        sse = SSE.new(response.stream)
        sse.write("{\"name\":\"John\"}")
        sse.write({ name: "Ryan" })
      ensure
        sse.close
      end

      def sse_with_event
        sse = SSE.new(response.stream, event: "send-name")
        sse.write("{\"name\":\"John\"}")
        sse.write({ name: "Ryan" })
      ensure
        sse.close
      end

      def sse_with_retry
        sse = SSE.new(response.stream, retry: 1000)
        sse.write("{\"name\":\"John\"}")
        sse.write({ name: "Ryan" }, retry: 1500)
      ensure
        sse.close
      end

      def sse_with_id
        sse = SSE.new(response.stream)
        sse.write("{\"name\":\"John\"}", id: 1)
        sse.write({ name: "Ryan" }, id: 2)
      ensure
        sse.close
      end

      def sse_with_multiple_line_message
        sse = SSE.new(response.stream)
        sse.write("first line.\nsecond line.")
      ensure
        sse.close
      end
    end

    tests SSETestController

    def wait_for_response_stream_close
      response.body
    end

    def test_basic_sse
      get :basic_sse

      wait_for_response_stream_close
      assert_match(/data: {\"name\":\"John\"}/, response.body)
      assert_match(/data: {\"name\":\"Ryan\"}/, response.body)
    end

    def test_sse_with_event_name
      get :sse_with_event

      wait_for_response_stream_close
      assert_match(/data: {\"name\":\"John\"}/, response.body)
      assert_match(/data: {\"name\":\"Ryan\"}/, response.body)
      assert_match(/event: send-name/, response.body)
    end

    def test_sse_with_retry
      get :sse_with_retry

      wait_for_response_stream_close
      first_response, second_response = response.body.split("\n\n")
      assert_match(/data: {\"name\":\"John\"}/, first_response)
      assert_match(/retry: 1000/, first_response)

      assert_match(/data: {\"name\":\"Ryan\"}/, second_response)
      assert_match(/retry: 1500/, second_response)
    end

    def test_sse_with_id
      get :sse_with_id

      wait_for_response_stream_close
      first_response, second_response = response.body.split("\n\n")
      assert_match(/data: {\"name\":\"John\"}/, first_response)
      assert_match(/id: 1/, first_response)

      assert_match(/data: {\"name\":\"Ryan\"}/, second_response)
      assert_match(/id: 2/, second_response)
    end

    def test_sse_with_multiple_line_message
      get :sse_with_multiple_line_message

      wait_for_response_stream_close
      first_response, second_response = response.body.split("\n")
      assert_match(/data: first line/, first_response)
      assert_match(/data: second line/, second_response)
    end
  end

  class LiveStreamTest < ActionController::TestCase
    class Exception < StandardError
    end

    class TestController < ActionController::Base
      include ActionController::Live

      attr_accessor :latch, :tc, :error_latch

      def self.controller_path
        'test'
      end

      def set_cookie
        cookies[:hello] = "world"
        response.stream.write "hello world"
        response.close
      end

      def render_text
        render plain: 'zomg'
      end

      def default_header
        response.stream.write "<html><body>hi</body></html>"
        response.stream.close
      end

      def basic_stream
        response.headers['Content-Type'] = 'text/event-stream'
        %w{ hello world }.each do |word|
          response.stream.write word
        end
        response.stream.close
      end

      def blocking_stream
        response.headers['Content-Type'] = 'text/event-stream'
        %w{ hello world }.each do |word|
          response.stream.write word
          latch.wait
        end
        response.stream.close
      end

      def thread_locals
        tc.assert_equal 'aaron', Thread.current[:setting]

        response.headers['Content-Type'] = 'text/event-stream'
        %w{ hello world }.each do |word|
          response.stream.write word
        end
        response.stream.close
      end

      def with_stale
        render plain: 'stale' if stale?(etag: "123", template: false)
      end

      def exception_in_view
        render 'doesntexist'
      end

      def exception_in_view_after_commit
        response.stream.write ""
        render 'doesntexist'
      end

      def exception_with_callback
        response.headers['Content-Type'] = 'text/event-stream'

        response.stream.on_error do
          response.stream.write %(data: "500 Internal Server Error"\n\n)
          response.stream.close
        end

        response.stream.write "" # make sure the response is committed
        raise 'An exception occurred...'
      end

      def exception_in_controller
        raise Exception, 'Exception in controller'
      end

      def bad_request_error
        raise ActionController::BadRequest
      end

      def exception_in_exception_callback
        response.headers['Content-Type'] = 'text/event-stream'
        response.stream.on_error do
          raise 'We need to go deeper.'
        end
        response.stream.write ''
        response.stream.write params[:widget][:didnt_check_for_nil]
      end

      def overfill_buffer_and_die
        logger = ActionController::Base.logger || Logger.new($stdout)
        response.stream.on_error do
          logger.warn 'Error while streaming'
          error_latch.count_down
        end

        # Write until the buffer is full. It doesn't expose that
        # information directly, so we must hard-code its size:
        10.times do
          response.stream.write '.'
        end
        # .. plus one more, because the #each frees up a slot:
        response.stream.write '.'

        latch.count_down

        # This write will block, and eventually raise
        response.stream.write 'x'

        20.times do
          response.stream.write '.'
        end
      end

      def ignore_client_disconnect
        response.stream.ignore_disconnect = true

        response.stream.write '' # commit

        # These writes will be ignored
        15.times do
          response.stream.write 'x'
        end

        logger.info 'Work complete'
        latch.count_down
      end
    end

    tests TestController

    def assert_stream_closed
      assert response.stream.closed?, 'stream should be closed'
      assert response.sent?, 'stream should be sent'
    end

    def capture_log_output
      output = StringIO.new
      old_logger, ActionController::Base.logger = ActionController::Base.logger, ActiveSupport::Logger.new(output)

      begin
        yield output
      ensure
        ActionController::Base.logger = old_logger
      end
    end

    def setup
      super

      def @controller.new_controller_thread
        Thread.new { yield }
      end
    end

    def test_set_cookie
      get :set_cookie
      assert_equal({'hello' => 'world'}, @response.cookies)
      assert_equal "hello world", @response.body
    end

    def test_write_to_stream
      get :basic_stream
      assert_equal "helloworld", @response.body
      assert_equal 'text/event-stream', @response.headers['Content-Type']
    end

    def test_async_stream
      rubinius_skip "https://github.com/rubinius/rubinius/issues/2934"

      @controller.latch = Concurrent::CountDownLatch.new
      parts             = ['hello', 'world']

      get :blocking_stream

      t = Thread.new(response) { |resp|
        resp.await_commit
        resp.stream.each do |part|
          assert_equal parts.shift, part
          ol = @controller.latch
          @controller.latch = Concurrent::CountDownLatch.new
          ol.count_down
        end
      }

      assert t.join(3), 'timeout expired before the thread terminated'
    end

    def test_abort_with_full_buffer
      @controller.latch = Concurrent::CountDownLatch.new
      @controller.error_latch = Concurrent::CountDownLatch.new

      capture_log_output do |output|
        get :overfill_buffer_and_die, :format => 'plain'

        t = Thread.new(response) { |resp|
          resp.await_commit
          _, _, body = resp.to_a
          body.each do
            @controller.latch.wait
            body.close
            break
          end
        }

        t.join
        @controller.error_latch.wait
        assert_match 'Error while streaming', output.rewind && output.read
      end
    end

    def test_ignore_client_disconnect
      @controller.latch = Concurrent::CountDownLatch.new

      capture_log_output do |output|
        get :ignore_client_disconnect

        t = Thread.new(response) { |resp|
          resp.await_commit
          _, _, body = resp.to_a
          body.each do
            body.close
            break
          end
        }

        t.join
        Timeout.timeout(3) do
          @controller.latch.wait
        end
        assert_match 'Work complete', output.rewind && output.read
      end
    end

    def test_thread_locals_get_copied
      @controller.tc = self
      Thread.current[:originating_thread] = Thread.current.object_id
      Thread.current[:setting]            = 'aaron'

      get :thread_locals
    end

    def test_live_stream_default_header
      get :default_header
      assert response.headers['Content-Type']
    end

    def test_render_text
      get :render_text
      assert_equal 'zomg', response.body
      assert_stream_closed
    end

    def test_exception_handling_html
      assert_raises(ActionView::MissingTemplate) do
        get :exception_in_view
      end

      capture_log_output do |output|
        get :exception_in_view_after_commit
        assert_match %r((window\.location = "/500\.html"</script></html>)$), response.body
        assert_match 'Missing template test/doesntexist', output.rewind && output.read
        assert_stream_closed
      end
      assert response.body
      assert_stream_closed
    end

    def test_exception_handling_plain_text
      assert_raises(ActionView::MissingTemplate) do
        get :exception_in_view, format: :json
      end

      capture_log_output do |output|
        get :exception_in_view_after_commit, format: :json
        assert_equal '', response.body
        assert_match 'Missing template test/doesntexist', output.rewind && output.read
        assert_stream_closed
      end
    end

    def test_exception_callback_when_committed
      current_threads = Thread.list

      capture_log_output do |output|
        get :exception_with_callback, format: 'text/event-stream'

        # Wait on the execution of all threads
        (Thread.list - current_threads).each(&:join)

        assert_equal %(data: "500 Internal Server Error"\n\n), response.body
        assert_match 'An exception occurred...', output.rewind && output.read
        assert_stream_closed
      end
    end

    def test_exception_in_controller_before_streaming
      assert_raises(ActionController::LiveStreamTest::Exception) do
        get :exception_in_controller, format: 'text/event-stream'
      end
    end

    def test_bad_request_in_controller_before_streaming
      assert_raises(ActionController::BadRequest) do
        get :bad_request_error, format: 'text/event-stream'
      end
    end

    def test_exceptions_raised_handling_exceptions_and_committed
      capture_log_output do |output|
        get :exception_in_exception_callback, format: 'text/event-stream'
        assert_equal '', response.body
        assert_match 'We need to go deeper', output.rewind && output.read
        assert_stream_closed
      end
    end

    def test_stale_without_etag
      get :with_stale
      assert_equal 200, response.status.to_i
    end

    def test_stale_with_etag
      @request.if_none_match = %(W/"#{Digest::MD5.hexdigest('123')}")
      get :with_stale
      assert_equal 304, response.status.to_i
    end
  end

  class BufferTest < ActionController::TestCase
    def test_nil_callback
      buf = ActionController::Live::Buffer.new nil
      assert buf.call_on_error
    end
  end
end

class LiveStreamRouterTest < ActionDispatch::IntegrationTest
  class TestController < ActionController::Base
    include ActionController::Live

    def index
      response.headers['Content-Type'] = 'text/event-stream'
      sse = SSE.new(response.stream)
      sse.write("{\"name\":\"John\"}")
      sse.write({ name: "Ryan" })
    ensure
      sse.close
    end
  end

  def self.call(env)
    routes.call(env)
  end

  def self.routes
    @routes ||= ActionDispatch::Routing::RouteSet.new
  end

  routes.draw do
    get '/test' => 'live_stream_router_test/test#index'
  end

  def app
    self.class
  end

  test "streaming served through the router" do
    get "/test"

    assert_response :ok
    assert_match(/data: {\"name\":\"John\"}/, response.body)
    assert_match(/data: {\"name\":\"Ryan\"}/, response.body)
  end
end