Add tests for `add_close` class

1 files changed, 50 insertions, 0 deletions

diff --git a/tests/test_wsgilib.py b/tests/test_wsgilib.py
new file mode 100644
index 0000000..f9366fc
--- /dev/null
+++ b/tests/test_wsgilib.py
@@ -0,0 +1,50 @@
+from paste.wsgilib import add_close
+
+
+def app_iterable_func_bytes():
+    yield b'a'
+    yield b'b'
+    yield b'c'
+
+
+def app_iterable_func_unicode():
+    yield b'a'.decode('ascii')
+    yield b'b'.decode('ascii')
+    yield b'c'.decode('ascii')
+
+
+def close_func():
+    global close_func_called
+    close_func_called = True
+
+
+def test_add_close_bytes():
+    global close_func_called
+
+    close_func_called = False
+    lst = []
+    app_iterable = app_iterable_func_bytes()
+
+    obj = add_close(app_iterable, close_func)
+    for x in obj:
+        lst.append(x)
+    obj.close()
+
+    assert lst == [b'a', b'b', b'c']
+    assert close_func_called
+
+
+def test_add_close_unicode():
+    global close_func_called
+
+    close_func_called = False
+    lst = []
+    app_iterable = app_iterable_func_unicode()
+
+    obj = add_close(app_iterable, close_func)
+    for x in obj:
+        lst.append(x)
+    obj.close()
+
+    assert lst == ['a', 'b', 'c']
+    assert close_func_called