• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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