From 1024af5285c2098f4284570caac64279a1aaa2a3 Mon Sep 17 00:00:00 2001
From: Miguel Grinberg <miguel.grinberg@gmail.com>
Date: Mon, 5 Jul 2021 00:27:22 +0100
Subject: [PATCH] More unit tests

---
 src/turbo_flask/turbo.py |  4 ++-
 tests/test_turbo.py      | 56 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/src/turbo_flask/turbo.py b/src/turbo_flask/turbo.py
index f3a7b41..feb09e3 100644
--- a/src/turbo_flask/turbo.py
+++ b/src/turbo_flask/turbo.py
@@ -174,9 +174,11 @@ def push(self, stream, to=None):
             to = self.clients.keys()
         elif not hasattr(to, '__len__') or isinstance(to, str):
             to = [to]
+        if hasattr(stream, '__len__') and not isinstance(stream, str):
+            stream = ''.join(stream)
         for recipient in to:
             for ws in self.clients[recipient]:
                 try:
                     ws.send(stream)
-                except ConnectionClosed:
+                except ConnectionClosed:  # pragma: no cover
                     pass
diff --git a/tests/test_turbo.py b/tests/test_turbo.py
index 9054475..9ebebe3 100644
--- a/tests/test_turbo.py
+++ b/tests/test_turbo.py
@@ -1,4 +1,5 @@
 import unittest
+from unittest import mock
 import pytest
 from flask import Flask, render_template_string
 from werkzeug.exceptions import NotFound
@@ -148,3 +149,58 @@ def test_streams(self):
             '<turbo-stream action="remove" target="bar">'
             '<template></template></turbo-stream>'
         )
+
+    def test_stream_response(self):
+        app = Flask(__name__)
+        turbo = turbo_flask.Turbo(app)
+
+        with app.test_request_context('/'):
+            r = turbo.stream([turbo.append('foo', 'bar'), turbo.remove('baz')])
+        assert r.get_data() == (
+            b'<turbo-stream action="append" target="bar">'
+            b'<template>foo</template>'
+            b'</turbo-stream>'
+            b'<turbo-stream action="remove" target="baz">'
+            b'<template></template>'
+            b'</turbo-stream>'
+        )
+
+    def test_push(self):
+        app = Flask(__name__)
+        turbo = turbo_flask.Turbo(app)
+        turbo.clients = {'123': [mock.MagicMock()], '456': [mock.MagicMock()]}
+
+        expected_stream = (
+            '<turbo-stream action="append" target="bar">'
+            '<template>foo</template>'
+            '</turbo-stream>'
+            '<turbo-stream action="remove" target="baz">'
+            '<template></template>'
+            '</turbo-stream>'
+        )
+        turbo.push([turbo.append('foo', 'bar'), turbo.remove('baz')])
+        turbo.clients['123'][0].send.assert_called_with(expected_stream)
+        turbo.clients['456'][0].send.assert_called_with(expected_stream)
+
+    def test_push_to(self):
+        app = Flask(__name__)
+        turbo = turbo_flask.Turbo(app)
+        turbo.clients = {'123': [mock.MagicMock()], '456': [mock.MagicMock()]}
+
+        expected_stream = (
+            '<turbo-stream action="append" target="bar">'
+            '<template>foo</template>'
+            '</turbo-stream>'
+            '<turbo-stream action="remove" target="baz">'
+            '<template></template>'
+            '</turbo-stream>'
+        )
+        turbo.push([turbo.append('foo', 'bar'), turbo.remove('baz')], to='456')
+        turbo.clients['123'][0].send.assert_not_called()
+        turbo.clients['456'][0].send.assert_called_with(expected_stream)
+        turbo.clients['123'][0].reset_mock()
+        turbo.clients['456'][0].reset_mock()
+        turbo.push([turbo.append('foo', 'bar'), turbo.remove('baz')],
+                   to=['123'])
+        turbo.clients['123'][0].send.assert_called_with(expected_stream)
+        turbo.clients['456'][0].send.assert_not_called()