TCPTransport.send returns Riemanns response

UDPTransport.send returns NotImplemented

Author: Sam Clements

riemann_client/client.py

@@ -39,8 +39,7 @@ def send_event(self, event):
         """Wraps an event in a message and sends it to Riemann"""
         message = riemann_client.riemann_pb2.Msg()
         message.events.extend([event])
-        self.transport.send(message)
-        return event
+        return self.transport.send(message)
 
     def event(self, **data):
         """Sends an event"""
@@ -65,8 +64,7 @@ def create_dict(event):
     def send_query(self, query):
         message = riemann_client.riemann_pb2.Msg()
         message.query.string = query
-        self.transport.send(message)
-        return self.transport.recv()
+        return self.transport.send(message)
 
     def query(self, query):
         response = self.send_query(query)

riemann_client/test_riemann_client.py

@@ -16,9 +16,6 @@ def connect(self):
 
     def send(self, message):
         self.string.write(message.SerializeToString())
-
-    def recv(self):
-        message = riemann_client.riemann_pb2.Msg()
         message.ok = True
         return message
 
@@ -38,27 +35,33 @@ def disconnect():
     return client
 
 
+@py.test.fixture
+def event(client):
+    """Send and return a blank event"""
+    message = client.event()
+    return message.events[0]
+
+
 class TestClient(object):
     def test_service(self, client):
         client.event(service='test event')
         assert 'test event' in client.transport.string.getvalue()
 
-    def test_default_hostname(self, client):
-        event = client.event()
+    def test_default_hostname(self, client, event):
         assert socket.gethostname() == event.host
         assert socket.gethostname() in client.transport.string.getvalue()
 
     def test_custom_hostname(self, client):
-        event = client.event(host="test.example.com")
+        event = client.event(host="test.example.com").events[0]
         assert "test.example.com" == event.host
         assert "test.example.com" in client.transport.string.getvalue()
 
     def test_tags(self, client):
         client.event(tags=['tag-1', 'tag-2'])
         assert "tag-1" in client.transport.string.getvalue()
 
-    def test_event_cls(self, client):
-        assert isinstance(client.event(), riemann_client.riemann_pb2.Event)
+    def test_event_cls(self, event):
+        assert isinstance(event, riemann_client.riemann_pb2.Event)
 
     def test_query(self, client):
         assert client.query("true") == []

riemann_client/transport.py

@@ -44,10 +44,6 @@ def disconnect(self):
     def send(self):
         pass
 
-    @abc.abstractmethod
-    def recv(self):
-        pass
-
 
 class UDPTransport(Transport):
     def connect(self):
@@ -58,9 +54,7 @@ def disconnect(self):
 
     def send(self, message):
         self.socket.sendto(message.SerializeToString(), self.address)
-
-    def recv(self):
-        raise NotImplementedError
+        return NotImplemented
 
 
 class TCPTransport(Transport):
@@ -75,18 +69,18 @@ def send(self, message):
         message = message.SerializeToString()
         self.socket.sendall(struct.pack('!I', len(message)) + message)
 
-    def recv(self):
         length = struct.unpack('!I', self.socket.recv(4))[0]
-
         response = riemann_client.riemann_pb2.Msg()
-        response.ParseFromString(self.recvall(length))
+        response.ParseFromString(self.socket_recvall(socket, length))
 
         if not response.ok:
             raise RiemannError(response.error)
+
         return response
 
-    def recvall(self, length, bufsize=4096):
+    @staticmethod
+    def socket_recvall(socket, length, bufsize=4096):
         data = ""
         while len(data) < length:
-            data += self.socket.recv(bufsize)
+            data += socket.recv(bufsize)
         return data