github.com/muhammedhassanm/blockchain@v0.0.0-20200120143007-697261defd4d/sawtooth-core-master/sdk/python/tests/test_zmq_driver.py (about)

     1  # Copyright 2018 Intel Corporation
     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  # -----------------------------------------------------------------------------
    15  
    16  import logging
    17  import threading
    18  import random
    19  import string
    20  import unittest
    21  import queue
    22  
    23  import zmq
    24  
    25  from sawtooth_sdk.consensus.engine import Engine
    26  from sawtooth_sdk.consensus.zmq_driver import ZmqDriver
    27  from sawtooth_sdk.protobuf import consensus_pb2
    28  from sawtooth_sdk.protobuf.validator_pb2 import Message
    29  
    30  
    31  LOGGER = logging.getLogger(__name__)
    32  
    33  
    34  class MockEngine(Engine):
    35      def __init__(self):
    36          self.updates = []
    37          self.exit = False
    38  
    39      def start(self, updates, service, startup_state):
    40          while not self.exit:
    41              try:
    42                  update = updates.get(timeout=1)
    43              except queue.Empty:
    44                  pass
    45              else:
    46                  self.updates.append(update)
    47  
    48      def stop(self):
    49          self.exit = True
    50  
    51      def name(self):
    52          return 'test-name'
    53  
    54      def version(self):
    55          return 'test-version'
    56  
    57  
    58  class TestDriver(unittest.TestCase):
    59      def setUp(self):
    60          self.ctx = zmq.Context.instance()
    61          self.socket = self.ctx.socket(zmq.ROUTER)
    62          self.socket.bind('tcp://127.0.0.1:*')
    63          self.url = self.socket.getsockopt_string(zmq.LAST_ENDPOINT)
    64          self.connection_id = None
    65  
    66          self.engine = MockEngine()
    67          self.driver = ZmqDriver(self.engine)
    68  
    69      def tearDown(self):
    70          self.socket.close()
    71  
    72      def recv_rep(self, request_type, response, response_type):
    73          # pylint: disable=unbalanced-tuple-unpacking
    74          connection_id, message_bytes = self.socket.recv_multipart(0)
    75  
    76          self.connection_id = connection_id
    77  
    78          message = Message()
    79          message.ParseFromString(message_bytes)
    80  
    81          request = request_type()
    82          request.ParseFromString(message.content)
    83  
    84          reply = Message(
    85              message_type=response_type,
    86              content=response.SerializeToString(),
    87              correlation_id=message.correlation_id)
    88  
    89          self.socket.send_multipart(
    90              [self.connection_id, reply.SerializeToString()],
    91              0)
    92  
    93          return request
    94  
    95      def send_req_rep(self, request, request_type):
    96          message = Message(
    97              message_type=request_type,
    98              correlation_id=generate_correlation_id(),
    99              content=request.SerializeToString())
   100  
   101          self.socket.send_multipart(
   102              [self.connection_id, message.SerializeToString()],
   103              0)
   104  
   105          # pylint: disable=unbalanced-tuple-unpacking
   106          _, reply_bytes = self.socket.recv_multipart(0)
   107  
   108          reply = Message()
   109          reply.ParseFromString(reply_bytes)
   110  
   111          self.assertEqual(
   112              reply.message_type,
   113              Message.CONSENSUS_NOTIFY_ACK)
   114  
   115          return reply.content
   116  
   117      def test_driver(self):
   118          # Start the driver in a separate thread to simulate the
   119          # validator and the driver
   120          driver_thread = threading.Thread(
   121              target=self.driver.start,
   122              args=(self.url,))
   123  
   124          driver_thread.start()
   125  
   126          response = consensus_pb2.ConsensusRegisterResponse(
   127              status=consensus_pb2.ConsensusRegisterResponse.OK)
   128  
   129          request = self.recv_rep(
   130              consensus_pb2.ConsensusRegisterRequest,
   131              response,
   132              Message.CONSENSUS_REGISTER_RESPONSE)
   133  
   134          self.assertEqual(request.name, 'test-name')
   135          self.assertEqual(request.version, 'test-version')
   136  
   137          self.send_req_rep(
   138              consensus_pb2.ConsensusNotifyPeerConnected(),
   139              Message.CONSENSUS_NOTIFY_PEER_CONNECTED)
   140  
   141          self.send_req_rep(
   142              consensus_pb2.ConsensusNotifyPeerDisconnected(),
   143              Message.CONSENSUS_NOTIFY_PEER_DISCONNECTED)
   144  
   145          self.send_req_rep(
   146              consensus_pb2.ConsensusNotifyPeerMessage(),
   147              Message.CONSENSUS_NOTIFY_PEER_MESSAGE)
   148  
   149          self.send_req_rep(
   150              consensus_pb2.ConsensusNotifyBlockNew(),
   151              Message.CONSENSUS_NOTIFY_BLOCK_NEW)
   152  
   153          self.send_req_rep(
   154              consensus_pb2.ConsensusNotifyBlockValid(),
   155              Message.CONSENSUS_NOTIFY_BLOCK_VALID)
   156  
   157          self.send_req_rep(
   158              consensus_pb2.ConsensusNotifyBlockInvalid(),
   159              Message.CONSENSUS_NOTIFY_BLOCK_INVALID)
   160  
   161          self.send_req_rep(
   162              consensus_pb2.ConsensusNotifyBlockCommit(),
   163              Message.CONSENSUS_NOTIFY_BLOCK_COMMIT)
   164  
   165          self.assertEqual(
   166              [msg_type for (msg_type, data) in self.engine.updates],
   167              [
   168                  Message.CONSENSUS_NOTIFY_PEER_CONNECTED,
   169                  Message.CONSENSUS_NOTIFY_PEER_DISCONNECTED,
   170                  Message.CONSENSUS_NOTIFY_PEER_MESSAGE,
   171                  Message.CONSENSUS_NOTIFY_BLOCK_NEW,
   172                  Message.CONSENSUS_NOTIFY_BLOCK_VALID,
   173                  Message.CONSENSUS_NOTIFY_BLOCK_INVALID,
   174                  Message.CONSENSUS_NOTIFY_BLOCK_COMMIT,
   175              ])
   176  
   177          self.driver.stop()
   178          driver_thread.join()
   179  
   180  
   181  def generate_correlation_id():
   182      return ''.join(random.choice(string.ascii_letters) for _ in range(16))