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  )
31  MOCK_RESP_WITHOUT_CALLBACK = (
32      b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1}'
33  )
34  MOCK_RESP_TEMPLATE = (
35      '{"id": %d, "result": 123, "error": null, "status": 1, "uid": 1, '
36      '"callback": null}'
37  )
38  MOCK_RESP_UNKNOWN_STATUS = (
39      b'{"id": 0, "result": 123, "error": null, "status": 0, "callback": null}'
40  )
41  MOCK_RESP_WITH_CALLBACK = (
42      b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1, '
43      b'"callback": "1-0"}'
44  )
45  MOCK_RESP_WITH_ERROR = b'{"id": 0, "error": 1, "status": 1, "uid": 1}'
46  MOCK_RESP_FLEXIABLE_RESULT_LENGTH = (
47      '{"id": 0, "result": "%s", "error": null, "status": 0, "callback": null}'
48  )
49
50  class MockSocketFile:
51
52    def __init__(self, resp):
53      self.resp = resp
54      self.last_write = None
55
56    def write(self, msg):
57      self.last_write = msg
58
59    def readline(self):
60      return self.resp
61
62    def flush(self):
63      pass
64
65  def setup_mock_socket_file(self, mock_create_connection, resp=MOCK_RESP):
66    """Sets up a fake socket file from the mock connection.
67
68    Args:
69      mock_create_connection: The mock method for creating a method.
70      resp: (str) response to give. MOCK_RESP by default.
71
72    Returns:
73      The mock file that will be injected into the code.
74    """
75    fake_file = self.MockSocketFile(resp)
76    fake_conn = mock.MagicMock()
77    fake_conn.makefile.return_value = fake_file
78    mock_create_connection.return_value = fake_conn
79    return fake_file
80
81  def generate_rpc_response(self, response_length=1024):
82    length = response_length - len(self.MOCK_RESP_FLEXIABLE_RESULT_LENGTH) + 2
83    chars = string.ascii_letters + string.digits
84    random_msg = ''.join(random.choice(chars) for i in range(length))
85    mock_response = self.MOCK_RESP_FLEXIABLE_RESULT_LENGTH % random_msg
86    return bytes(mock_response, 'utf-8')
87