Mocked telemetry in modularinput tests

Author: amysutedja

splunklib/modularinput/script.py

@@ -13,15 +13,19 @@
 # under the License.
 
 from __future__ import absolute_import
+
+import sys
+
 from abc import ABCMeta, abstractmethod
+import splunklib
+from splunklib import six
 from splunklib.six.moves.urllib.parse import urlsplit
-import sys
 
 from ..client import Service
 from .event_writer import EventWriter
 from .input_definition import InputDefinition
 from .validation_definition import ValidationDefinition
-from splunklib import six
+from ..wire._internal import Telemetry, TelemetryMetric
 
 try:
     import xml.etree.cElementTree as ET
@@ -70,6 +74,20 @@ def run_script(self, args, event_writer, input_stream):
                 # passed on stdin as XML, and the script will write events on
                 # stdout and log entries on stderr.
                 self._input_definition = InputDefinition.parse(input_stream)
+
+                # create a telemetry metric
+                metric = TelemetryMetric(**{
+                    'metric_type': 'event',
+                    'component': 'splunk-sdk-python',
+                    'data': {
+                        'version': splunklib.__version__
+                    }
+                })
+
+                # call out to telemetry
+                telemetry = Telemetry(self.service)
+                telemetry.submit(metric.to_wire())
+
                 self.stream_events(self._input_definition, event_writer)
                 event_writer.close()
                 return 0

splunklib/wire/_internal/json_sink.py

@@ -44,6 +44,7 @@ def submit(self, data):
         """
 
         response = self._post('', headers=self.__class__.JSON_HEADER, body=json.dumps(data))
+
         body = json.loads(response.body.read().decode('utf-8'))
 
         return response, body

tests/modularinput/test_script.py

@@ -1,8 +1,12 @@
+import io
+import json
 import sys
 
+from mock import Mock, patch
+
+from splunklib import six
 from splunklib.client import Service
 from splunklib.modularinput import Script, EventWriter, Scheme, Argument, Event
-import io
 
 from splunklib.modularinput.utils import xml_compare
 from tests.modularinput.modularinput_testlib import data_open
@@ -14,6 +18,17 @@
 
 TEST_SCRIPT_PATH = "__IGNORED_SCRIPT_PATH__"
 
+PATCHED_TELEMETRY_RESPONSE = {
+    'status': 201,
+    'reason': 'Created',
+    'body.read.return_value': six.ensure_binary(json.dumps({
+        'message': 'Data submitted successfully',
+        'metricValueID': '26844DB9-7806-40E0-96C0-1BD554930BA8'
+    })),
+    'headers': [
+        ('content-type', 'application/json; charset=UTF-8')
+    ]
+}
 
 def test_error_on_script_with_null_scheme(capsys):
     """A script that returns a null scheme should generate no output on
@@ -184,9 +199,12 @@ def stream_events(self, inputs, ew):
     script = NewScript()
     input_configuration = data_open("data/conf_with_2_inputs.xml")
 
-    ew = EventWriter(sys.stdout, sys.stderr)
+    event_writer = EventWriter(sys.stdout, sys.stderr)
 
-    return_value = script.run_script([TEST_SCRIPT_PATH], ew, input_configuration)
+    with patch.object(Service, 'post') as patched_telemetry_post:
+        patched_telemetry_post.return_value = Mock(**PATCHED_TELEMETRY_RESPONSE)
+
+        return_value = script.run_script([TEST_SCRIPT_PATH], event_writer, input_configuration)
 
     output = capsys.readouterr()
     assert output.err == ""
@@ -218,7 +236,12 @@ def stream_events(self, inputs, ew):
             self.authority_uri = inputs.metadata['server_uri']
 
     script = NewScript()
-    with data_open("data/conf_with_2_inputs.xml") as input_configuration:
+
+    with data_open("data/conf_with_2_inputs.xml") as input_configuration, \
+            patch.object(Service, 'post') as patched_telemetry_post:
+
+        patched_telemetry_post.return_value = Mock(**PATCHED_TELEMETRY_RESPONSE)
+
         ew = EventWriter(sys.stdout, sys.stderr)
 
         assert script.service is None

tests/test_telemetry.py

@@ -21,15 +21,10 @@
 from splunklib.wire._internal.telemetry import Telemetry
 from splunklib.wire._internal.telemetry_metric import TelemetryMetric
 
-try:
-    import unittest
-except ImportError:
-    import unittest2 as unittest
-
 @pytest.mark.app
-class TelemetryTestCase(testlib.SDKTestCase):
+class TestTelemetry(testlib.SDKTestCase):
     def setUp(self):
-        super(TelemetryTestCase, self).setUp()
+        super(TestTelemetry, self).setUp()
 
         self.service.namespace['owner'] = 'nobody'
         self.service.namespace['app'] = 'sdk-app-collection'
@@ -51,10 +46,3 @@ def test_submit(self):
 
         # it should return a 201
         self.assertEqual(response.status, 201)
-
-if __name__ == "__main__":
-    try:
-        import unittest2 as unittest
-    except ImportError:
-        import unittest
-    unittest.main()

tox.ini

@@ -27,7 +27,8 @@ passenv = LANG
 setenv = SPLUNK_HOME=/opt/splunk
          INPUT_EXAMPLE_UPLOAD=/opt/splunk/var/log/splunk/splunkd_ui_access.log
 whitelist_externals = true
-deps = pytest
+deps = mock
+       pytest
        pytest-cov
        xmlrunner
        unittest2