test_zmq_driver.py

# Copyright 2018 Intel Corporation
#
# 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.
# -----------------------------------------------------------------------------

import logging
import threading
import random
import string
import unittest
import queue

import zmq

from sawtooth_sdk.consensus.engine import Engine
from sawtooth_sdk.consensus.zmq_driver import ZmqDriver
from sawtooth_sdk.protobuf import consensus_pb2
from sawtooth_sdk.protobuf import network_pb2
from sawtooth_sdk.protobuf.validator_pb2 import Message


LOGGER = logging.getLogger(__name__)


class MockEngine(Engine):
    # Ignore invalid override pylint issues
    # pylint: disable=invalid-overridden-method
    def __init__(self):
        self.updates = []
        self.exit = False

    def start(self, updates, service, startup_state):
        while not self.exit:
            try:
                update = updates.get(timeout=1)
            except queue.Empty:
                pass
            else:
                self.updates.append(update)

    def stop(self):
        self.exit = True

    # Ignore invalid override pylint issues
    # pylint: disable=invalid-overridden-method
    def name(self):
        return 'test-name'

    # Ignore invalid override pylint issues
    # pylint: disable=invalid-overridden-method
    def version(self):
        return 'test-version'

    def additional_protocols(self):
        return [('Test-Name', 'Test-Version')]


class TestDriver(unittest.TestCase):
    def setUp(self):
        self.ctx = zmq.Context.instance()
        self.socket = self.ctx.socket(zmq.ROUTER)
        self.socket.bind('tcp://127.0.0.1:*')
        self.url = self.socket.getsockopt_string(zmq.LAST_ENDPOINT)
        self.connection_id = None

        self.engine = MockEngine()
        self.driver = ZmqDriver(self.engine)

    def tearDown(self):
        self.socket.close()

    def recv_rep(self, request_type, response, response_type):
        # pylint: disable=unbalanced-tuple-unpacking
        connection_id, message_bytes = self.socket.recv_multipart(0)

        self.connection_id = connection_id

        message = Message()
        message.ParseFromString(message_bytes)

        request = request_type()
        request.ParseFromString(message.content)

        reply = Message(
            message_type=response_type,
            content=response.SerializeToString(),
            correlation_id=message.correlation_id)

        self.socket.send_multipart(
            [self.connection_id, reply.SerializeToString()],
            0)

        return request

    def send_req_rep(self, request, request_type):
        message = Message(
            message_type=request_type,
            correlation_id=generate_correlation_id(),
            content=request.SerializeToString())

        self.socket.send_multipart(
            [self.connection_id, message.SerializeToString()],
            0)

        # pylint: disable=unbalanced-tuple-unpacking
        _, reply_bytes = self.socket.recv_multipart(0)

        reply = Message()
        reply.ParseFromString(reply_bytes)

        self.assertEqual(
            reply.message_type,
            Message.CONSENSUS_NOTIFY_ACK)

        return reply.content

    def test_driver(self):
        # Start the driver in a separate thread to simulate the
        # validator and the driver
        driver_thread = threading.Thread(
            target=self.driver.start,
            args=(self.url,))

        driver_thread.start()

        response = consensus_pb2.ConsensusRegisterResponse(
            status=consensus_pb2.ConsensusRegisterResponse.OK)

        request = self.recv_rep(
            consensus_pb2.ConsensusRegisterRequest,
            response,
            Message.CONSENSUS_REGISTER_RESPONSE)

        additional_protocols = \
            [(p.name, p.version) for p in request.additional_protocols]

        self.assertEqual(request.name, 'test-name')
        self.assertEqual(request.version, 'test-version')
        self.assertEqual(additional_protocols, [('Test-Name', 'Test-Version')])

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyEngineActivated(),
            Message.CONSENSUS_NOTIFY_ENGINE_ACTIVATED)

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyPeerConnected(),
            Message.CONSENSUS_NOTIFY_PEER_CONNECTED)

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyPeerDisconnected(),
            Message.CONSENSUS_NOTIFY_PEER_DISCONNECTED)

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyPeerMessage(),
            Message.CONSENSUS_NOTIFY_PEER_MESSAGE)

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyBlockNew(),
            Message.CONSENSUS_NOTIFY_BLOCK_NEW)

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyBlockValid(),
            Message.CONSENSUS_NOTIFY_BLOCK_VALID)

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyBlockInvalid(),
            Message.CONSENSUS_NOTIFY_BLOCK_INVALID)

        self.send_req_rep(
            consensus_pb2.ConsensusNotifyBlockCommit(),
            Message.CONSENSUS_NOTIFY_BLOCK_COMMIT)

        self.send_req_rep(
            network_pb2.PingRequest(),
            Message.PING_REQUEST)

        self.assertEqual(
            [msg_type for (msg_type, data) in self.engine.updates],
            [
                Message.CONSENSUS_NOTIFY_PEER_CONNECTED,
                Message.CONSENSUS_NOTIFY_PEER_DISCONNECTED,
                Message.CONSENSUS_NOTIFY_PEER_MESSAGE,
                Message.CONSENSUS_NOTIFY_BLOCK_NEW,
                Message.CONSENSUS_NOTIFY_BLOCK_VALID,
                Message.CONSENSUS_NOTIFY_BLOCK_INVALID,
                Message.CONSENSUS_NOTIFY_BLOCK_COMMIT,
            ])

        self.driver.stop()
        driver_thread.join()


def generate_correlation_id():
    return ''.join(random.choice(string.ascii_letters) for _ in range(16))