1# Copyright 2017 Google Inc. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14 15import random 16import string 17import unittest 18from unittest import mock 19 20 21class JsonRpcClientTestBase(unittest.TestCase): 22 """Base class for tests of JSONRPC clients. 23 24 Contains infrastructure for mocking responses. 25 """ 26 27 MOCK_RESP = ( 28 b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1, ' 29 b'"callback": null}') 30 MOCK_RESP_WITHOUT_CALLBACK = ( 31 b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1}') 32 MOCK_RESP_TEMPLATE = ( 33 '{"id": %d, "result": 123, "error": null, "status": 1, "uid": 1, ' 34 '"callback": null}') 35 MOCK_RESP_UNKNOWN_STATUS = ( 36 b'{"id": 0, "result": 123, "error": null, "status": 0, ' 37 b'"callback": null}') 38 MOCK_RESP_WITH_CALLBACK = ( 39 b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1, ' 40 b'"callback": "1-0"}') 41 MOCK_RESP_WITH_ERROR = b'{"id": 0, "error": 1, "status": 1, "uid": 1}' 42 MOCK_RESP_FLEXIABLE_RESULT_LENGTH = ( 43 '{"id": 0, "result": "%s", "error": null, "status": 0, "callback": null}') 44 45 class MockSocketFile: 46 47 def __init__(self, resp): 48 self.resp = resp 49 self.last_write = None 50 51 def write(self, msg): 52 self.last_write = msg 53 54 def readline(self): 55 return self.resp 56 57 def flush(self): 58 pass 59 60 def setup_mock_socket_file(self, mock_create_connection, resp=MOCK_RESP): 61 """Sets up a fake socket file from the mock connection. 62 63 Args: 64 mock_create_connection: The mock method for creating a method. 65 resp: (str) response to give. MOCK_RESP by default. 66 67 Returns: 68 The mock file that will be injected into the code. 69 """ 70 fake_file = self.MockSocketFile(resp) 71 fake_conn = mock.MagicMock() 72 fake_conn.makefile.return_value = fake_file 73 mock_create_connection.return_value = fake_conn 74 return fake_file 75 76 def generate_rpc_response(self, response_length=1024): 77 length = response_length - len(self.MOCK_RESP_FLEXIABLE_RESULT_LENGTH) + 2 78 chars = string.ascii_letters + string.digits 79 random_msg = ''.join(random.choice(chars) for i in range(length)) 80 mock_response = self.MOCK_RESP_FLEXIABLE_RESULT_LENGTH % random_msg 81 return bytes(mock_response, 'utf-8') 82