stateobject: use typing, enable tuples and more complex datatypes

2 files changed, 102 insertions, 52 deletions

diff --git a/test/mitmproxy/test_stateobject.py b/test/mitmproxy/test_stateobject.py
index d8c7a8e9..bd5d1792 100644
--- a/test/mitmproxy/test_stateobject.py
+++ b/test/mitmproxy/test_stateobject.py
@@ -1,101 +1,146 @@
-from typing import List
+import typing
+
 import pytest
 
 from mitmproxy.stateobject import StateObject
 
 
-class Child(StateObject):
+class TObject(StateObject):
     def __init__(self, x):
         self.x = x
 
-    _stateobject_attributes = dict(
-        x=int
-    )
-
     @classmethod
     def from_state(cls, state):
         obj = cls(None)
         obj.set_state(state)
         return obj
 
+
+class Child(TObject):
+    _stateobject_attributes = dict(
+        x=int
+    )
+
     def __eq__(self, other):
         return isinstance(other, Child) and self.x == other.x
 
 
-class Container(StateObject):
-    def __init__(self):
-        self.child = None
-        self.children = None
-        self.dictionary = None
+class TTuple(TObject):
+    _stateobject_attributes = dict(
+        x=typing.Tuple[int, Child]
+    )
+
+
+class TList(TObject):
+    _stateobject_attributes = dict(
+        x=typing.List[Child]
+    )
+
 
+class TDict(TObject):
     _stateobject_attributes = dict(
-        child=Child,
-        children=List[Child],
-        dictionary=dict,
+        x=typing.Dict[str, Child]
     )
 
-    @classmethod
-    def from_state(cls, state):
-        obj = cls()
-        obj.set_state(state)
-        return obj
+
+class TAny(TObject):
+    _stateobject_attributes = dict(
+        x=typing.Any
+    )
+
+
+class TSerializableChild(TObject):
+    _stateobject_attributes = dict(
+        x=Child
+    )
 
 
 def test_simple():
     a = Child(42)
+    assert a.get_state() == {"x": 42}
     b = a.copy()
-    assert b.get_state() == {"x": 42}
     a.set_state({"x": 44})
     assert a.x == 44
     assert b.x == 42
 
 
-def test_container():
-    a = Container()
-    a.child = Child(42)
+def test_serializable_child():
+    child = Child(42)
+    a = TSerializableChild(child)
+    assert a.get_state() == {
+        "x": {"x": 42}
+    }
+    a.set_state({
+        "x": {"x": 43}
+    })
+    assert a.x.x == 43
+    assert a.x is child
     b = a.copy()
-    assert a.child.x == b.child.x
-    b.child.x = 44
-    assert a.child.x != b.child.x
+    assert a.x == b.x
+    assert a.x is not b.x
 
 
-def test_container_list():
-    a = Container()
-    a.children = [Child(42), Child(44)]
+def test_tuple():
+    a = TTuple((42, Child(43)))
     assert a.get_state() == {
-        "child": None,
-        "children": [{"x": 42}, {"x": 44}],
-        "dictionary": None,
+        "x": (42, {"x": 43})
     }
-    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())
+    b = a.copy()
+    a.set_state({"x": (44, {"x": 45})})
+    assert a.x == (44, Child(45))
+    assert b.x == (42, Child(43))
+
 
+def test_tuple_err():
+    a = TTuple(None)
+    with pytest.raises(ValueError, msg="Invalid data"):
+        a.set_state({"x": (42,)})
 
-def test_container_dict():
-    a = Container()
-    a.dictionary = dict()
-    a.dictionary['foo'] = 'bar'
-    a.dictionary['bar'] = Child(44)
+
+def test_list():
+    a = TList([Child(1), Child(2)])
     assert a.get_state() == {
-        "child": None,
-        "children": None,
-        "dictionary": {'bar': {'x': 44}, 'foo': 'bar'},
+        "x": [{"x": 1}, {"x": 2}],
     }
     copy = a.copy()
-    assert len(copy.dictionary) == 2
-    assert copy.dictionary is not a.dictionary
-    assert copy.dictionary['bar'] is not a.dictionary['bar']
+    assert len(copy.x) == 2
+    assert copy.x is not a.x
+    assert copy.x[0] is not a.x[0]
+
+
+def test_dict():
+    a = TDict({"foo": Child(42)})
+    assert a.get_state() == {
+        "x": {"foo": {"x": 42}}
+    }
+    b = a.copy()
+    assert list(a.x.items()) == list(b.x.items())
+    assert a.x is not b.x
+    assert a.x["foo"] is not b.x["foo"]
+
+
+def test_any():
+    a = TAny(42)
+    b = a.copy()
+    assert a.x == b.x
+
+    a = TAny(object())
+    with pytest.raises(AssertionError):
+        a.get_state()
 
 
 def test_too_much_state():
-    a = Container()
-    a.child = Child(42)
+    a = Child(42)
     s = a.get_state()
     s['foo'] = 'bar'
-    b = Container()
 
     with pytest.raises(RuntimeWarning):
-        b.set_state(s)
+        a.set_state(s)
+
+
+def test_none():
+    a = Child(None)
+    assert a.get_state() == {"x": None}
+    a = Child(42)
+    a.set_state({"x": None})
+    assert a.x is None
diff --git a/test/mitmproxy/utils/test_typecheck.py b/test/mitmproxy/utils/test_typecheck.py
index 5295fff5..9cb4334e 100644
--- a/test/mitmproxy/utils/test_typecheck.py
+++ b/test/mitmproxy/utils/test_typecheck.py
@@ -93,3 +93,8 @@ def test_typesec_to_str():
     assert(typecheck.typespec_to_str(typing.Optional[str])) == "optional str"
     with pytest.raises(NotImplementedError):
         typecheck.typespec_to_str(dict)
+
+
+def test_mapping_types():
+    # this is not covered by check_option_type, but still belongs in this module
+    assert (str, int) == typecheck.mapping_types(typing.Mapping[str, int])