Merge pull request #2114 from mitmproxy/fix-websocket-serialization

make websocket flows serializable

2 files changed, 34 insertions, 3 deletions

diff --git a/test/mitmproxy/test_stateobject.py b/test/mitmproxy/test_stateobject.py
index 7b8e30d0..d8c7a8e9 100644
--- a/test/mitmproxy/test_stateobject.py
+++ b/test/mitmproxy/test_stateobject.py
@@ -26,10 +26,12 @@ class Container(StateObject):
     def __init__(self):
         self.child = None
         self.children = None
+        self.dictionary = None
 
     _stateobject_attributes = dict(
         child=Child,
         children=List[Child],
+        dictionary=dict,
     )
 
     @classmethod
@@ -62,12 +64,30 @@ def test_container_list():
     a.children = [Child(42), Child(44)]
     assert a.get_state() == {
         "child": None,
-        "children": [{"x": 42}, {"x": 44}]
+        "children": [{"x": 42}, {"x": 44}],
+        "dictionary": None,
     }
     copy = a.copy()
     assert len(copy.children) == 2
     assert copy.children is not a.children
     assert copy.children[0] is not a.children[0]
+    assert Container.from_state(a.get_state())
+
+
+def test_container_dict():
+    a = Container()
+    a.dictionary = dict()
+    a.dictionary['foo'] = 'bar'
+    a.dictionary['bar'] = Child(44)
+    assert a.get_state() == {
+        "child": None,
+        "children": None,
+        "dictionary": {'bar': {'x': 44}, 'foo': 'bar'},
+    }
+    copy = a.copy()
+    assert len(copy.dictionary) == 2
+    assert copy.dictionary is not a.dictionary
+    assert copy.dictionary['bar'] is not a.dictionary['bar']
 
 
 def test_too_much_state():
diff --git a/test/mitmproxy/test_websocket.py b/test/mitmproxy/test_websocket.py
index f2963390..62f69e2d 100644
--- a/test/mitmproxy/test_websocket.py
+++ b/test/mitmproxy/test_websocket.py
@@ -1,5 +1,7 @@
+import io
 import pytest
 
+from mitmproxy.contrib import tnetstring
 from mitmproxy import flowfilter
 from mitmproxy.test import tflow
 
@@ -14,8 +16,6 @@ class TestWebSocketFlow:
         b = f2.get_state()
         del a["id"]
         del b["id"]
-        del a["handshake_flow"]["id"]
-        del b["handshake_flow"]["id"]
         assert a == b
         assert not f == f2
         assert f is not f2
@@ -60,3 +60,14 @@ class TestWebSocketFlow:
         assert 'WebSocketFlow' in repr(f)
         assert 'binary message: ' in repr(f.messages[0])
         assert 'text message: ' in repr(f.messages[1])
+
+    def test_serialize(self):
+        b = io.BytesIO()
+        d = tflow.twebsocketflow().get_state()
+        tnetstring.dump(d, b)
+        assert b.getvalue()
+
+        b = io.BytesIO()
+        d = tflow.twebsocketflow().handshake_flow.get_state()
+        tnetstring.dump(d, b)
+        assert b.getvalue()