test the hec obj a little more

Author: jettero

hubblestack/hec/obj.py

@@ -18,6 +18,7 @@
 
 from . dq import DiskQueue, NoQueue, QueueCapacityError
 from hubblestack.utils.stdrec import update_payload
+from hubblestack.utils.encoding import encode_something_to_bytes
 
 __version__ = '1.0'
 
@@ -259,7 +260,7 @@ def __init__(self, token, index, http_event_server, host='', http_event_port='80
             md5 = hashlib.md5()
             uril = sorted([ x.uri for x in self.server_uri ])
             for u in uril:
-                md5.update(u)
+                md5.update(encode_something_to_bytes(u))
             actual_disk_queue = os.path.join(disk_queue, md5.hexdigest())
             log.debug("disk_queue for %s: %s", uril, actual_disk_queue)
             self.queue = DiskQueue(actual_disk_queue, size=disk_queue_size, compression=disk_queue_compression)
@@ -299,12 +300,12 @@ def _queue_event(self, payload, meta_data=None):
             log.error("disk queue is full, dropping payload")
 
 
-    def queueEvent(self, dat, eventtime=''):
+    def queueEvent(self, dat, eventtime='', no_queue=False):
         if not isinstance(dat, Payload):
             dat = Payload(dat, eventtime, no_queue=no_queue)
         if dat.no_queue: # here you silly hec, queue this no_queue payload...
             return
-        count_input(payload)
+        count_input(dat)
         self._queue_event(dat)
 
     def flushQueue(self):

tests/unittests/test_hec_dq.py

@@ -1,3 +1,5 @@
+# coding: utf-8
+
 import pytest
 import os
 

tests/unittests/test_hec_obj.py

@@ -0,0 +1,33 @@
+# coding: utf-8
+
+import os
+import json
+import mock
+from hubblestack.hec import HEC
+
+TEST_DQ_DIR = os.environ.get('TEST_DQ_DIR', '/tmp/dq.{0}'.format(os.getuid()))
+
+@mock.patch.object(HEC, '_send')
+def test_hec__send_trivially(mock_send):
+    hec = HEC('token', 'index', 'server')
+    hec.sendEvent({'test': 'test-tacular'})
+    assert json.loads(mock_send.call_args.args[0].dat)['test'] == 'test-tacular'
+
+@mock.patch.object(HEC, '_send') # just in case, not actually used
+def test_queue_things_with_compression(mock_send, __opts__, __salt__):
+    hec = HEC('token', 'index', 'server',
+        disk_queue=TEST_DQ_DIR, disk_queue_size=1000,
+        disk_queue_compression=9)
+
+    results_of_side_effect = list()
+    def side_effect(x):
+        results_of_side_effect.append(x)
+    mock_send.side_effect = side_effect
+
+    gz = list()
+    for i in range(100):
+        dat = {f'event{i}': f'test{i}'}
+        hec.queueEvent(dat)
+        gz.append( json.dumps(dat) )
+    hec.flushQueue()
+    assert ' '.join(results_of_side_effect) == ' '.join(gz)