Class hierarchy for TelemetryMetrics

Author: amysutedja

splunklib/modularinput/script.py

@@ -25,7 +25,7 @@
 from .event_writer import EventWriter
 from .input_definition import InputDefinition
 from .validation_definition import ValidationDefinition
-from ..wire._internal import Telemetry, TelemetryMetric
+from ..wire._internal import Telemetry, EventTelemetryMetric
 
 try:
     import xml.etree.cElementTree as ET
@@ -76,8 +76,7 @@ def run_script(self, args, event_writer, input_stream):
                 self._input_definition = InputDefinition.parse(input_stream)
 
                 # create a telemetry metric
-                metric = TelemetryMetric(**{
-                    'metric_type': 'event',
+                metric = EventTelemetryMetric(**{
                     'component': 'splunk-sdk-python',
                     'data': {
                         'version': splunklib.__version__

splunklib/wire/_internal/__init__.py

@@ -14,5 +14,6 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+from .aggregate_telemetry_metric import *
+from .event_telemetry_metric import *
 from .telemetry import *
-from .telemetry_metric import *

splunklib/wire/_internal/aggregate_telemetry_metric.py

@@ -0,0 +1,69 @@
+# coding=utf-8
+#
+# Copyright © 2011-2020 Splunk, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"): you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+from splunklib.wire._internal.telemetry_metric import TelemetryMetric
+
+class AggregateTelemetryMetric(TelemetryMetric):
+    METRIC_TYPE = 'aggregate'
+
+    def __init__(self, component, data,
+                 opt_in_required=2,
+                 version=None,
+                 timestamp=None,
+                 visibility=None,
+                 index_data=None,
+                 begin=None,
+                 end=None):
+        super(AggregateTelemetryMetric, self).__init__(
+            AggregateTelemetryMetric.METRIC_TYPE,
+            component,
+            data,
+            opt_in_required=opt_in_required,
+            version=version,
+            timestamp=timestamp,
+            visibility=visibility,
+            index_data=index_data
+        )
+
+        self.begin = begin
+        self.end = end
+
+    @property
+    def begin(self):
+        return self._begin
+
+    @begin.setter
+    def begin(self, value):
+        self._begin = value
+
+    @property
+    def end(self):
+        return self._end
+
+    @end.setter
+    def end(self, value):
+        self._end = value
+
+    def to_wire(self):
+        wire = super(AggregateTelemetryMetric, self).to_wire()
+
+        if self.begin is not None:
+            wire['begin'] = self.begin
+
+        if self.end is not None:
+            wire['end'] = self.end
+
+        return wire

splunklib/wire/_internal/event_telemetry_metric.py

@@ -0,0 +1,69 @@
+# coding=utf-8
+#
+# Copyright © 2011-2020 Splunk, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"): you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+from splunklib.wire._internal.telemetry_metric import TelemetryMetric
+
+class EventTelemetryMetric(TelemetryMetric):
+    METRIC_TYPE = 'event'
+
+    def __init__(self, component, data,
+                 opt_in_required=2,
+                 version=None,
+                 timestamp=None,
+                 visibility=None,
+                 index_data=None,
+                 user_id=None,
+                 experience_id=None):
+        super(EventTelemetryMetric, self).__init__(
+            EventTelemetryMetric.METRIC_TYPE,
+            component,
+            data,
+            opt_in_required=opt_in_required,
+            version=version,
+            timestamp=timestamp,
+            visibility=visibility,
+            index_data=index_data
+        )
+
+        self.user_id = user_id
+        self.experience_id = experience_id
+
+    @property
+    def user_id(self):
+        return self._user_id
+
+    @user_id.setter
+    def user_id(self, value):
+        self._user_id = value
+
+    @property
+    def experience_id(self):
+        return self._experience_id
+
+    @experience_id.setter
+    def experience_id(self, value):
+        self._experience_id = value
+
+    def to_wire(self):
+        wire = super(EventTelemetryMetric, self).to_wire()
+
+        if self.user_id is not None:
+            wire['userID'] = self.user_id
+
+        if self.experience_id is not None:
+            wire['experienceID'] = self.experience_id
+
+        return wire

splunklib/wire/_internal/telemetry_metric.py

@@ -14,12 +14,24 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
-class TelemetryMetric:
-    def __init__(self, metric_type, component, data, opt_in_required=2):
+from abc import ABCMeta
+from splunklib import six
+
+class TelemetryMetric(six.with_metaclass(ABCMeta, object)):
+    def __init__(self, metric_type, component, data,
+                 opt_in_required=2,
+                 version=None,
+                 index_data=None,
+                 timestamp=None,
+                 visibility=None):
         self.metric_type = metric_type
         self.component = component
         self.data = data
         self.opt_in_required = opt_in_required
+        self.version = version
+        self.index_data = index_data
+        self.timestamp = timestamp
+        self.visibility = visibility
 
     @property
     def metric_type(self):
@@ -53,10 +65,56 @@ def opt_in_required(self):
     def opt_in_required(self, value):
         self._opt_in_required = value
 
+    @property
+    def version(self):
+        return self._version
+
+    @version.setter
+    def version(self, value):
+        self._version = value
+
+    @property
+    def index_data(self):
+        return self._index_data
+
+    @index_data.setter
+    def index_data(self, value):
+        self._index_data = value
+
+    @property
+    def timestamp(self):
+        return self._timestamp
+
+    @timestamp.setter
+    def timestamp(self, value):
+        self._timestamp = value
+
+    @property
+    def visibility(self):
+        return self._visibility
+
+    @visibility.setter
+    def visibility(self, value):
+        self._visibility = value
+
     def to_wire(self):
-        return {
+        wire = {
             'type': self.metric_type,
             'component': self.component,
             'data': self.data,
             'optInRequired': self.opt_in_required,
         }
+
+        if self.version is not None:
+            wire['version'] = self.version
+
+        if self.index_data is not None:
+            wire['indexData'] = self.index_data
+
+        if self.timestamp is not None:
+            wire['timestamp'] = self.timestamp
+
+        if self.visibility is not None:
+            wire['visibility'] = self.visibility
+
+        return wire

tests/modularinput/test_script.py

@@ -4,7 +4,7 @@
 
 from mock import Mock, patch
 
-from splunklib import six
+from splunklib import six, __version__
 from splunklib.client import Service
 from splunklib.modularinput import Script, EventWriter, Scheme, Argument, Event
 
@@ -219,6 +219,15 @@ def stream_events(self, inputs, ew):
 def test_telemetry(capsys):
     """Check that writing telemetry goes smoothly."""
 
+    EXPECTED_TELEMETRY_BODY = {
+        'type': 'event',
+        'component': 'splunk-sdk-python',
+        'data': {
+            'version': __version__,
+        },
+        'optInRequired': 2
+    }
+
     # Override abstract methods
     class NewScript(Script):
         def get_scheme(self):
@@ -246,7 +255,7 @@ def stream_events(self, _inputs, ew):
         assert post_args == ('telemetry-metric/',)
         assert post_kwargs == {
             'app': None,
-            'body': '{"type": "event", "component": "splunk-sdk-python", "data": {"version": "1.6.13"}, "optInRequired": 2}',
+            'body': json.dumps(EXPECTED_TELEMETRY_BODY),
             'headers': [('Content-Type', 'application/json')],
             'owner': None,
             'sharing': None

tests/test_telemetry.py

@@ -18,8 +18,7 @@
 import pytest
 
 from tests import testlib
-from splunklib.wire._internal.telemetry import Telemetry
-from splunklib.wire._internal.telemetry_metric import TelemetryMetric
+from splunklib.wire._internal import Telemetry, EventTelemetryMetric, AggregateTelemetryMetric
 
 @pytest.mark.app
 class TestTelemetry(testlib.SDKTestCase):
@@ -33,8 +32,7 @@ def setUp(self):
 
     def test_submit(self):
         # create a telemetry metric
-        metric = TelemetryMetric(**{
-            'metric_type': 'event',
+        metric = EventTelemetryMetric(**{
             'component': 'telemetry_test_case',
             'data': {
                 'testValue': 32
@@ -46,3 +44,43 @@ def test_submit(self):
 
         # it should return a 201
         self.assertEqual(response.status, 201)
+
+    def test_event_submit(self):
+        # create a telemetry metric
+        metric = EventTelemetryMetric(**{
+            'component': 'telemetry_test_case',
+            'data': {
+                'testValue': 32
+            },
+            'version': 'test',
+            'index_data': False,
+            'timestamp': 0,
+            'visibility': ['anonymous'],
+        })
+
+        # call out to telemetry
+        response, _body = self.telemetry.submit(metric.to_wire())
+
+        # it should return a 201
+        self.assertEqual(response.status, 201)
+
+    def test_aggregate_submit(self):
+        # create a telemetry metric
+        metric = AggregateTelemetryMetric(**{
+            'component': 'telemetry_test_case',
+            'data': {
+                'testValue': 32
+            },
+            'version': 'test',
+            'index_data': False,
+            'timestamp': 3,
+            'visibility': ['anonymous'],
+            'begin': 0,
+            'end': 1,
+        })
+
+        # call out to telemetry
+        response, _body = self.telemetry.submit(metric.to_wire())
+
+        # it should return a 201
+        self.assertEqual(response.status, 201)