1 files changed, 48 insertions, 38 deletions

diff --git a/tests/topotests/lib/mcast-tester.py b/tests/topotests/lib/mcast-tester.py
index 07e4ab8773..30beccb787 100755
--- a/tests/topotests/lib/mcast-tester.py
+++ b/tests/topotests/lib/mcast-tester.py
@@ -21,25 +21,25 @@ for the multicast group we subscribed to.
 """
 
 import argparse
-import os
 import json
+import os
 import socket
-import subprocess
 import struct
+import subprocess
 import sys
 import time
 
+
 #
 # Functions
 #
 def interface_name_to_index(name):
     "Gets the interface index using its name. Returns None on failure."
-    interfaces = json.loads(
-        subprocess.check_output('ip -j link show', shell=True))
+    interfaces = json.loads(subprocess.check_output("ip -j link show", shell=True))
 
     for interface in interfaces:
-        if interface['ifname'] == name:
-            return interface['ifindex']
+        if interface["ifname"] == name:
+            return interface["ifindex"]
 
     return None
 
@@ -59,13 +59,12 @@ def multicast_join(sock, ifindex, group, port):
 # Main code.
 #
 parser = argparse.ArgumentParser(description="Multicast RX utility")
-parser.add_argument('socket', help='Point to topotest UNIX socket')
-parser.add_argument('group', help='Multicast IP')
-parser.add_argument('interface', help='Interface name')
+parser.add_argument("group", help="Multicast IP")
+parser.add_argument("interface", help="Interface name")
+parser.add_argument("--socket", help="Point to topotest UNIX socket")
 parser.add_argument(
-    '--send',
-    help='Transmit instead of join with interval (defaults to 0.7 sec)',
-    type=float, default=0)
+    "--send", help="Transmit instead of join with interval", type=float, default=0
+)
 args = parser.parse_args()
 
 ttl = 16
@@ -74,7 +73,7 @@ port = 1000
 # Get interface index/validate.
 ifindex = interface_name_to_index(args.interface)
 if ifindex is None:
-    sys.stderr.write('Interface {} does not exists\n'.format(args.interface))
+    sys.stderr.write("Interface {} does not exists\n".format(args.interface))
     sys.exit(1)
 
 # We need root privileges to set up multicast.
@@ -83,47 +82,58 @@ if os.geteuid() != 0:
     sys.exit(1)
 
 # Wait for topotest to synchronize with us.
-toposock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, 0)
-while True:
-    try:
-        toposock.connect(args.socket)
-        break
-    except ConnectionRefusedError:
-        time.sleep(1)
-        continue
+if not args.socket:
+    toposock = None
+else:
+    toposock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM, 0)
+    while True:
+        try:
+            toposock.connect(args.socket)
+            break
+        except ConnectionRefusedError:
+            time.sleep(1)
+            continue
+    # Set topotest socket non blocking so we can multiplex the main loop.
+    toposock.setblocking(False)
 
 msock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
 if args.send > 0:
     # Prepare multicast bit in that interface.
     msock.setsockopt(
-        socket.SOL_SOCKET, 25,
-        struct.pack("%ds" % len(args.interface),
-                    args.interface.encode('utf-8')))
+        socket.SOL_SOCKET,
+        25,
+        struct.pack("%ds" % len(args.interface), args.interface.encode("utf-8")),
+    )
     # Set packets TTL.
-    msock.setsockopt(
-        socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, struct.pack("b", ttl))
+    msock.setsockopt(socket.IPPROTO_IP, socket.IP_MULTICAST_TTL, struct.pack("b", ttl))
     # Block to ensure packet send.
     msock.setblocking(True)
-    # Set topotest socket non blocking so we can multiplex the main loop.
-    toposock.setblocking(False)
 else:
     multicast_join(msock, ifindex, args.group, port)
 
+
+def should_exit():
+    if not toposock:
+        # If we are sending then we have slept
+        if not args.send:
+            time.sleep(100)
+        return False
+    else:
+        try:
+            data = toposock.recv(1)
+            if data == b"":
+                print(" -> Connection closed")
+                return True
+        except BlockingIOError:
+            return False
+
+
 counter = 0
-while True:
+while not should_exit():
     if args.send > 0:
         msock.sendto(b"test %d" % counter, (args.group, port))
         counter += 1
         time.sleep(args.send)
 
-    try:
-        data = toposock.recv(1)
-        if data == b'':
-            print(' -> Connection closed')
-            break
-    except BlockingIOError:
-        continue
-
 msock.close()
-
 sys.exit(0)