Make send_directive more flexible

Author: jimi-c

receptor/controller.py

@@ -9,19 +9,18 @@
 logger = logging.getLogger(__name__)
 
 
-def send_directive(directive, recipient, payload, socket_path, expected_responses=1):
-    if payload == '-':
-        payload = sys.stdin.read()
+def connect_to_socket(socket_path):
     sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
     sock.connect(socket_path)
+    return sock
+
+def send_directive(directive, recipient, payload, sock):
+    if payload == '-':
+        payload = sys.stdin.read()
     sock.sendall(f"{recipient}\n{directive}\n{payload}".encode('utf-8') + protocol.DELIM)
     response = b''
-    response_count = 0
-    while response_count < expected_responses:
-        response = sock.recv(4096)
-        sys.stdout.buffer.write(response + b"\n")
-        sys.stdout.flush()
-        response_count += 1
+    response = sock.recv(4096)
+    return response
 
 # FIXME: the socket path is in the config, it shouldn't need to be passed as an arg here
 def mainloop(receptor, socket_path, loop=asyncio.get_event_loop()):

receptor/entrypoints.py

@@ -1,5 +1,7 @@
 import datetime
+import json
 import logging
+import sys
 import time
 
 from .receptor import Receptor
@@ -17,20 +19,41 @@ def run_as_controller(config):
 
 def run_as_ping(config):
     logger.info(f'Sending ping to {config.ping_recipient}.')
-    now = datetime.datetime.utcnow()
-    pings_sent = 0
-    while True:
-        controller.send_directive('receptor:ping', config.ping_recipient, now.isoformat(), config.ping_socket_path)
-        pings_sent += 1
-        if config.ping_count != 0 and pings_sent >= config.ping_count:
-            break
-        elif config.ping_delay > 0.0:
-            time.sleep(config.ping_delay)
+    sock = controller.connect_to_socket(config.ping_socket_path)
+
+    try:
+        pings_sent = 0
+        while True:
+            now = datetime.datetime.utcnow()
+            response = controller.send_directive('receptor:ping', config.ping_recipient, now.isoformat(), sock)
+            resp_json = json.loads(response)
+            if 'code' in resp_json and resp_json['code'] != 0:
+                sys.stdout.buffer.write(b"Failed to ping node.\n")
+            else:
+                sys.stdout.buffer.write(response + b"\n")
+            sys.stdout.flush()
+            pings_sent += 1
+            if config.ping_count != 0 and pings_sent >= config.ping_count:
+                break
+            elif config.ping_delay > 0.0:
+                time.sleep(config.ping_delay)
+    except KeyboardInterrupt:
+        pass
+    finally:
+        sock.close()
 
 
 def run_as_send(config):
     logger.info(f'Sending a {config.send_directive} directive to {config.send_recipient}.')
-    controller.send_directive(config.send_directive, config.send_recipient, config.send_payload, config.send_socket_path)
+    sock = controller.connect_to_socket(config.send_socket_path)
+    try:
+        response = controller.send_directive(config.send_directive, config.send_recipient, config.send_payload, sock)
+        sys.stdout.buffer.write(response + b"\n")
+        sys.stdout.flush()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        sock.close()
 
 
 def run_as_node(config):