diff options
| -rw-r--r-- | tests/test_core.py | 21 | ||||
| -rw-r--r-- | viewstate/core.py | 21 |
2 files changed, 21 insertions, 21 deletions
diff --git a/tests/test_core.py b/tests/test_core.py index bfb6c4f..2a3b11c 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -20,34 +20,28 @@ class TestViewState(object): def test_invalid_decode(self): with pytest.raises(ViewStateException): - vs = ViewState() - vs.raw = b'\x01\x02' + vs = ViewState(raw=b'\x01\x02') vs.decode() def test_parse_const_value(self): - vs = ViewState() - vs.raw = b'\xff\x01\x67' + vs = ViewState(raw=b'\xff\x01\x67') assert vs.decode() is True def test_parse_string_value(self): - vs = ViewState() s = 'abcdefghij' - vs.raw = b'\xff\x01\x05\x0a' + s.encode() + vs = ViewState(raw=b'\xff\x01\x05\x0a' + s.encode()) assert vs.decode() == s def test_parse_simple_dict(self): - vs = ViewState() - vs.raw = b'\xff\x01\x18\x02\x05\x01a\x05\x01b\x05\x01c\x05\x01d' + vs = ViewState(raw=b'\xff\x01\x18\x02\x05\x01a\x05\x01b\x05\x01c\x05\x01d') assert vs.decode() == {'a': 'b', 'c': 'd'} def test_parse_simple_list(self): - vs = ViewState() - vs.raw = b'\xff\x01\x16\x05\x05\x01a\x05\x01b\x05\x01c\x05\x01d\x05\x01e' + vs = ViewState(raw=b'\xff\x01\x16\x05\x05\x01a\x05\x01b\x05\x01c\x05\x01d\x05\x01e') assert vs.decode() == ['a', 'b', 'c', 'd', 'e'] def test_parse_simple_pair(self): - vs = ViewState() - vs.raw = b'\xff\x01\x0f\x67\x68' + vs = ViewState(raw=b'\xff\x01\x0f\x67\x68') assert vs.decode() == (True, False) def test_parse_complex_pair(self): @@ -61,6 +55,5 @@ class TestViewState(object): def test_parse_unknown(self): with pytest.raises(ViewStateException): - vs = ViewState() - vs.raw = b'\xff\x01\x99\x99\x99' + vs = ViewState(raw=b'\xff\x01\x99\x99\x99') assert vs.decode() diff --git a/viewstate/core.py b/viewstate/core.py index d9567c0..b86afc7 100644 --- a/viewstate/core.py +++ b/viewstate/core.py @@ -67,13 +67,16 @@ def parse(b): class ViewState(object): - def __init__(self, base64=''): - self.base64 = base64 + def __init__(self, base64=None, raw=None): + if base64: + self.base64 = base64 + try: + self.raw = b64decode(self.base64) + except BinAsciiError as bae: + raise ViewStateException('Cannot decode base64 input') + elif raw: + self.raw = raw self.decoded = None - try: - self.raw = b64decode(self.base64) - except BinAsciiError as bae: - raise ViewStateException('Cannot decode base64 input') @property def preamble(self): @@ -87,7 +90,11 @@ class ViewState(object): format_marker = b'\xff' version_marker = b'\x01' preamble = format_marker + version_marker - return self.preamble == preamble + + try: + return self.preamble == preamble + except AttributeError: + return False def decode(self): if not self.is_valid(): |