Merge pull request #48 from twitter/test-quorum-packets

tests(quorum_packet): test Request, Proposal, Ack, Commit & Ping

zktraffic/network/sniffer.py

@@ -44,7 +44,7 @@ class Sniffer(Thread):
   """
   class RegistrationError(Exception): pass
 
-  def __init__(self, iface, port, msg_cls, handler=None, dump_bad_packet=False):
+  def __init__(self, iface, port, msg_cls, handler=None, dump_bad_packet=False, start=True):
     super(Sniffer, self).__init__()
     self.setDaemon(True)
 
@@ -58,7 +58,8 @@ def __init__(self, iface, port, msg_cls, handler=None, dump_bad_packet=False):
     if handler is not None:
       self.add_handler(handler)
 
-    self.start()
+    if start:  # pragma: no cover
+      self.start()
 
   def add_handler(self, handler):
     if handler is None:
@@ -69,17 +70,22 @@ def add_handler(self, handler):
 
     self._handlers.append(handler)
 
-  def run(self):
+  def run(self, *args, **kwargs):
     pfilter = "port %d" % self._port
     try:
-      if self._iface == "any":
-        sniff(filter=pfilter, store=0, prn=self.handle_packet)
-      else:
-        sniff(filter=pfilter, store=0, prn=self.handle_packet, iface=self._iface)
+      sniff_kwargs = {"filter": pfilter, "store": 0, "prn": self.handle_packet}
+      if self._iface != "any":
+        sniff_kwargs["iface"] = self._iface
+
+      if "offline" in kwargs:
+        sniff_kwargs["offline"] = kwargs["offline"]
+
+      sniff(**sniff_kwargs)
     except socket.error as ex:
       sys.stderr.write("Error: %s, device: %s\n" % (ex, self._iface))
     finally:
-      os.kill(os.getpid(), signal.SIGINT)
+      if "offline" not in kwargs:
+        os.kill(os.getpid(), signal.SIGINT)
 
   def handle_packet(self, packet):
     try:

zktraffic/tests/common.py

@@ -26,6 +26,9 @@
 )
 
 
-def consume_packets(capture_file, zkt):
-  path = os.path.join(_resources_dir, "%s.pcap" % (capture_file))
-  sniff(offline=path, prn=zkt.handle_packet)
+def get_full_path(name):
+  return os.path.join(_resources_dir, "%s.pcap" % (name))
+
+
+def consume_packets(capture_file, sniffer):
+  sniff(offline=get_full_path(capture_file), prn=sniffer.handle_packet)

zktraffic/tests/test_zab.py

@@ -17,11 +17,23 @@
 import unittest
 
 from zktraffic.base.network import BadPacket
+from zktraffic.base.zookeeper import OpCodes
+from zktraffic.network.sniffer import Sniffer
 from zktraffic.zab.quorum_packet import (
+  Ack,
+  Commit,
   PacketType,
-  QuorumPacket
+  Ping,
+  Proposal,
+  QuorumPacket,
+  Request,
 )
 
+from .common import get_full_path
+
+
+LEADER_PORT = 20022
+
 
 class ZabTestCase(unittest.TestCase):
   def test_basic(self):
@@ -33,3 +45,78 @@ def test_basic(self):
     packet = QuorumPacket.from_payload(payload, '127.0.0.1:2889', '127.0.0.1:10000', 0)
     self.assertEqual(PacketType.PROPOSAL, packet.type)
     self.assertEqual(packet.zxid, 2000)
+
+  def test_from_pcap(self):
+    requests = []
+    proposals = []
+    commits = []
+    acks = []
+    pings = []
+
+    def handler(message):
+      if isinstance(message, Request):
+        requests.append(message)
+      elif isinstance(message, Proposal):
+        proposals.append(message)
+      elif isinstance(message, Commit):
+        commits.append(message)
+      elif isinstance(message, Ack):
+        acks.append(message)
+      elif isinstance(message, Ping):
+        pings.append(message)
+
+    sniffer = Sniffer(
+      iface="test",
+      port=LEADER_PORT,
+      msg_cls=QuorumPacket,
+      handler=handler,
+      dump_bad_packet=False,
+      start=False
+    )
+
+    sniffer.run(offline=get_full_path('zab_request'))
+
+    # requests
+    assert len(requests) == 3
+
+    assert requests[0].req_type == OpCodes.CREATESESSION
+    assert requests[0].session_id_literal == "0x1001df405af0000"
+    assert requests[0].zxid == -1
+
+    assert requests[1].req_type == OpCodes.CREATE
+    assert requests[1].session_id_literal == "0x1001df405af0000"
+    assert requests[0].zxid == -1
+
+    assert requests[1].req_type == OpCodes.CREATE
+    assert requests[1].session_id_literal == "0x1001df405af0000"
+    assert requests[0].zxid == -1
+
+    # proposals
+    assert len(proposals) == 6  # 2 createSession + 4 create
+
+    assert proposals[0].txn_type == OpCodes.CREATESESSION
+    assert proposals[0].session_id_literal == "0x1001df405af0000"
+    assert proposals[0].zxid_literal == "0x100000001"
+
+    assert proposals[2].txn_type == OpCodes.CREATE
+    assert proposals[2].session_id_literal == "0x1001df405af0000"
+    assert proposals[2].zxid_literal == "0x100000002"
+
+    # commits
+    assert len(commits) == 6  # 2 createSession + 4 create
+
+    assert commits[0].zxid_literal == "0x100000001"
+    assert commits[2].zxid_literal == "0x100000002"
+    assert commits[4].zxid_literal == "0x100000003"
+
+    # acks
+    assert len(acks) == 6  # 2 createSession + 4 create
+
+    assert acks[0].zxid_literal == "0x100000001"
+    assert acks[2].zxid_literal == "0x100000002"
+    assert acks[4].zxid_literal == "0x100000003"
+
+    # pings
+    assert len(pings) == 57  # for such a short run, this numbers looks too high
+
+    assert pings[0].zxid_literal == "0x100000000"