websockets: refactor implementation and add tests

1 files changed, 23 insertions, 0 deletions

diff --git a/test/netlib/websockets/test_masker.py b/test/netlib/websockets/test_masker.py
new file mode 100644
index 00000000..528fce71
--- /dev/null
+++ b/test/netlib/websockets/test_masker.py
@@ -0,0 +1,23 @@
+import codecs
+import pytest
+
+from netlib import websockets
+
+
+class TestMasker(object):
+
+    @pytest.mark.parametrize("input,expected", [
+        ([b"a"], '00'),
+        ([b"four"], '070d1616'),
+        ([b"fourf"], '070d161607'),
+        ([b"fourfive"], '070d1616070b1501'),
+        ([b"a", b"aasdfasdfa", b"asdf"], '000302170504021705040205120605'),
+        ([b"a" * 50, b"aasdfasdfa", b"asdf"], '00030205000302050003020500030205000302050003020500030205000302050003020500030205000302050003020500030205120605051206050500110702'),  # noqa
+    ])
+    def test_masker(self, input, expected):
+        m = websockets.Masker(b"abcd")
+        data = b"".join([m(t) for t in input])
+        assert data == codecs.decode(expected, 'hex')
+
+        data = websockets.Masker(b"abcd")(data)
+        assert data == b"".join(input)