tests/lib/jsonrpc_client_test_base.py

# Copyright 2017 Google Inc.
# 
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# 
#     http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from builtins import str

import mock
import unittest


class JsonRpcClientTestBase(unittest.TestCase):
    """Base class for tests of JSONRPC clients.

    Contains infrastructure for mocking responses.
    """

    MOCK_RESP = (
        b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1, '
        b'"callback": null}')
    MOCK_RESP_WITHOUT_CALLBACK = (
        b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1}')
    MOCK_RESP_TEMPLATE = (
        '{"id": %d, "result": 123, "error": null, "status": 1, "uid": 1, '
        '"callback": null}')
    MOCK_RESP_UNKNOWN_STATUS = (
        b'{"id": 0, "result": 123, "error": null, "status": 0, '
        b'"callback": null}')
    MOCK_RESP_WITH_CALLBACK = (
        b'{"id": 0, "result": 123, "error": null, "status": 1, "uid": 1, '
        b'"callback": "1-0"}')
    MOCK_RESP_WITH_ERROR = b'{"id": 0, "error": 1, "status": 1, "uid": 1}'

    class MockSocketFile(object):
        def __init__(self, resp):
            self.resp = resp
            self.last_write = None

        def write(self, msg):
            self.last_write = msg

        def readline(self):
            return self.resp

        def flush(self):
            pass

    def setup_mock_socket_file(self, mock_create_connection, resp=MOCK_RESP):
        """Sets up a fake socket file from the mock connection.

        Args:
            mock_create_connection: The mock method for creating a method.
            resp: (str) response to give. MOCK_RESP by default.

        Returns:
            The mock file that will be injected into the code.
        """
        fake_file = self.MockSocketFile(resp)
        fake_conn = mock.MagicMock()
        fake_conn.makefile.return_value = fake_file
        mock_create_connection.return_value = fake_conn
        return fake_file