diff --git a/tests/test_handler.py b/tests/test_handler.py index 2f30d9d..a9ad924 100644 --- a/tests/test_handler.py +++ b/tests/test_handler.py @@ -1,3 +1,4 @@ +import io import unittest import paramiko @@ -7,7 +8,7 @@ from tests.utils import read_file, make_tests_data_path from webssh import handler from webssh import worker from webssh.handler import ( - MixinHandler, WsockHandler, PrivateKey, InvalidValueError + IndexHandler, MixinHandler, WsockHandler, PrivateKey, InvalidValueError, SSHClient ) try: @@ -315,3 +316,25 @@ class TestWsockHandler(unittest.TestCase): obj.worker_ref = ref WsockHandler.on_message(obj, b'{"data": "somestuff"}') obj.close.assert_called_with(reason='Worker closed') + +class TestIndexHandler(unittest.TestCase): + def test_null_in_encoding(self): + handler = Mock(spec=IndexHandler) + + # This is a little nasty, but the index handler has a lot of + # dependencies to mock. Mocking out everything but the bits + # we want to test lets us test this case without needing to + # refactor the relevant code out of IndexHandler + def parse_encoding(data): + return IndexHandler.parse_encoding(handler, data) + handler.parse_encoding = parse_encoding + + ssh = Mock(spec=SSHClient) + stdin = io.BytesIO() + stdout = io.BytesIO(initial_bytes=b"UTF-8\0") + stderr = io.BytesIO() + ssh.exec_command.return_value = (stdin, stdout, stderr) + + encoding = IndexHandler.get_default_encoding(handler, ssh) + self.assertEquals("utf-8", encoding) + diff --git a/webssh/utils.py b/webssh/utils.py index 845ca56..454b1cf 100644 --- a/webssh/utils.py +++ b/webssh/utils.py @@ -59,6 +59,8 @@ def is_valid_encoding(encoding): u'test'.encode(encoding) except LookupError: return False + except ValueError: + return False return True