Added support for Riemann event attributes (closes #2)

Includes many more tests, and which have been moved to a new directory

Author: Sam Clements

riemann_client/client.py

@@ -29,6 +29,11 @@ def create_event(data):
         event = riemann_client.riemann_pb2.Event()
         event.host = socket.gethostname()
         event.tags.extend(data.pop('tags', []))
+
+        for key, value in data.pop('attributes', {}).items():
+            attribute = event.attributes.add()
+            attribute.key, attribute.value = key, value
+
         for name, value in data.items():
             if value is not None:
                 setattr(event, name, value)
@@ -55,6 +60,7 @@ def create_dict(event):
             'service': event.service,
             'tags': list(event.tags),
             'ttl': event.ttl,
+            'attributes': dict(((a.key, a.value) for a in event.attributes)),
             'metric_f': event.metric_f,
             'metric_d': event.metric_d,
             'metric_sint64': event.metric_sint64

riemann_client/test_riemann_client.py

@@ -0,0 +1,25 @@
+from __future__ import absolute_import, unicode_literals
+
+import StringIO
+
+import py.test
+
+import riemann_client.transport
+
+
+class StringTransport(riemann_client.transport.Transport):
+    def connect(self):
+        self.string = StringIO.StringIO()
+
+    def send(self, message):
+        self.string.write(message.SerializeToString())
+        message.ok = True
+        return message
+
+    def disconnect(self):
+        self.string.close()
+
+
+@py.test.fixture
+def string_transport():
+    return StringTransport()

riemann_client/tests/conftest.py

@@ -0,0 +1,114 @@
+from __future__ import absolute_import, unicode_literals
+
+import socket
+import uuid
+
+import py.test
+
+import riemann_client.client
+import riemann_client.riemann_pb2
+import riemann_client.transport
+
+
+@py.test.fixture
+def client(request, string_transport):
+    client = riemann_client.client.Client(transport=string_transport)
+    client.transport.connect()
+
+    @request.addfinalizer
+    def disconnect():
+        client.transport.disconnect()
+
+    return client
+
+
+@py.test.fixture
+def unique():
+    return str(uuid.uuid4())
+
+
+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):
+        event = client.event().events[0]
+        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").events[0]
+        assert "test.example.com" == event.host
+        assert "test.example.com" in client.transport.string.getvalue()
+
+    def test_tags(self, client, unique):
+        client.event(tags=[unique])
+        assert unique in client.transport.string.getvalue()
+
+    def test_attributes(self, client, unique):
+        client.event(attributes={'key': unique})
+        assert unique in client.transport.string.getvalue()
+
+    def test_event_cls(self, event):
+        assert isinstance(event, riemann_client.riemann_pb2.Event)
+
+    def test_query(self, client):
+        assert client.query('true') == []
+
+    def test_udp_query(self):
+        transport = riemann_client.transport.UDPTransport()
+        client = riemann_client.client.Client(transport)
+
+        with py.test.raises(Exception):
+            client.query('true')
+
+
+@py.test.fixture
+def event(unique):
+    return riemann_client.client.Client.create_event({
+        'host': 'test.example.com',
+        'tags': [unique],
+        'attributes': {
+            unique: unique
+        }
+    })
+
+
+class TestCreateEvent(object):
+    def test_host(self, event):
+        assert event.host == 'test.example.com'
+
+    def test_tags(self, event, unique):
+        assert unique in event.tags
+
+    def test_attributes(self, event, unique):
+        assert event.attributes[0].key == unique
+        assert event.attributes[0].value == unique
+
+
+@py.test.fixture
+def event_as_dict(event):
+    return riemann_client.client.Client.create_dict(event)
+
+
+class TestCreateDict(object):
+    def test_host(self, event, event_as_dict):
+        assert event.host == event_as_dict['host']
+
+    def test_host_type(self, event_as_dict):
+        assert isinstance(event_as_dict['host'], basestring)
+
+    def test_tags(self, event, event_as_dict):
+        for tag in event.tags:
+            assert tag in event_as_dict['tags']
+
+    def test_tags_type(self, event_as_dict):
+        assert isinstance(event_as_dict['tags'], list)
+
+    def test_attibutes(self, event, event_as_dict):
+        for attr in event.attributes:
+            assert event_as_dict['attributes'][attr.key] == attr.value
+
+    def test_attibutes_type(self, event_as_dict):
+        assert isinstance(event_as_dict['attributes'], dict)

riemann_client/tests/test_riemann_client.py

@@ -0,0 +1,67 @@
+from __future__ import absolute_import, unicode_literals
+
+import py.test
+
+import riemann_client.client
+import riemann_client.riemann_pb2
+import riemann_client.transport
+
+
+@py.test.fixture
+def queued_client(request, string_transport):
+    """A Riemann client using the StringIO transport and QueuedClient"""
+    client = riemann_client.client.QueuedClient(transport=string_transport)
+    client.transport.connect()
+
+    @request.addfinalizer
+    def disconnect():
+        client.transport.disconnect()
+
+    return client
+
+
+@py.test.fixture
+def using_simple_queue(queued_client):
+    """An event queue with a single event"""
+    queued_client.event(service='test')
+
+
+@py.test.fixture
+def large_queue(queued_client):
+    """An event queue with 100 events"""
+    items = ['-->{0}<--'.format(i) for i in range(0, 1000)]
+    for description in items:
+        queued_client.event(service='queue', description=description)
+    return items
+
+
+def test_simple_queue_event(queued_client, using_simple_queue):
+    assert queued_client.queue.events[0].service == 'test'
+
+
+def test_simple_queue_length(queued_client, using_simple_queue):
+    assert len(queued_client.queue.events) == 1
+
+
+def test_simple_queue_event_not_sent(queued_client, using_simple_queue):
+    assert "test" not in queued_client.transport.string.getvalue()
+
+
+def test_simple_queue_event_sent(queued_client, using_simple_queue):
+    queued_client.flush()
+    assert "test" in queued_client.transport.string.getvalue()
+
+
+def test_deciqueue_length(queued_client, large_queue):
+    assert len(queued_client.queue.events) == len(large_queue)
+
+
+def test_deciqueue_output(queued_client, large_queue):
+    queued_client.flush()
+    for description in large_queue:
+        assert description in queued_client.transport.string.getvalue()
+
+
+def test_deciqueue_flush(queued_client, large_queue):
+    queued_client.flush()
+    assert len(queued_client.queue.events) == 0

riemann_client/tests/test_riemann_queued_client.py

@@ -1,6 +1,6 @@
 [tox]
 minversion=1.6.0
-envlist=py26,py27,py26-flake8
+envlist=py26,py27,py26-flake8,py26-coverage
 
 [testenv]
 commands=py.test riemann_client
@@ -43,7 +43,7 @@ omit=
 exclude_lines=
     def __repr__
     raise NotImplementedError
-    class NullHandler
+    pass
 
 [html]
 title=riemann-client coverage report