Merge pull request #319 from splunk/use_output_redirect_for_modinput_tests

shilpabijam · web-flow · commit 8bcebf7366b1 · 2020-05-14T17:14:00.000-07:00
removed forcing to bytes or text, test stdout/err with pytest fixtures
diff --git a/splunklib/modularinput/event_writer.py b/splunklib/modularinput/event_writer.py
@@ -16,7 +16,7 @@
 import sys
 
 from io import TextIOWrapper, TextIOBase
-from splunklib.six import ensure_text
+from splunklib.six import ensure_str
 from .event import ET
 
 try:
@@ -43,15 +43,8 @@ def __init__(self, output = sys.stdout, error = sys.stderr):
         :param output: Where to write the output; defaults to sys.stdout.
         :param error: Where to write any errors; defaults to sys.stderr.
         """
-        if isinstance(output, TextIOBase):
-            self._out = output
-        else:
-            self._out = TextIOWrapper(output)
-
-        if isinstance(error, TextIOBase):
-            self._err = error
-        else:
-            self._err = TextIOWrapper(error)
+        self._out = output
+        self._err = error
 
         # has the opening <stream> tag been written yet?
         self.header_written = False
@@ -63,18 +56,20 @@ def write_event(self, event):
         """
 
         if not self.header_written:
-            self._out.write(ensure_text("<stream>"))
+            self._out.write("<stream>")
             self.header_written = True
 
         event.write_to(self._out)
 
     def log(self, severity, message):
         """Logs messages about the state of this modular input to Splunk.
         These messages will show up in Splunk's internal logs.
+
         :param severity: ``string``, severity of message, see severities defined as class constants.
         :param message: ``string``, message to log.
         """
-        self._err.write(ensure_text("%s %s\n" % (severity, message)))
+
+        self._err.write("%s %s\n" % (severity, message))
         self._err.flush()
 
     def write_xml_document(self, document):
@@ -83,11 +78,10 @@ def write_xml_document(self, document):
 
         :param document: An ``ElementTree`` object.
         """
-        data = ET.tostring(document)
-        self._out.write(ensure_text(data))
+        self._out.write(ensure_str(ET.tostring(document)))
         self._out.flush()
 
     def close(self):
         """Write the closing </stream> tag to make this XML well formed."""
-        self._out.write(ensure_text("</stream>"))
+        self._out.write("</stream>")
         self._out.flush()
diff --git a/splunklib/modularinput/script.py b/splunklib/modularinput/script.py
@@ -105,8 +105,7 @@ def run_script(self, args, event_writer, input_stream):
                 return 1
 
         except Exception as e:
-            err_string = EventWriter.ERROR + str(e)
-            event_writer._err.write(err_string)
+            event_writer.log(EventWriter.ERROR, str(e))
             return 1
 
     @property
diff --git a/tests/modularinput/test_event.py b/tests/modularinput/test_event.py
@@ -15,143 +15,139 @@
 # under the License.
 
 from __future__ import absolute_import
+
+import sys
+
+import pytest
+
 from tests.modularinput.modularinput_testlib import unittest, xml_compare, data_open
 from splunklib.modularinput.event import Event, ET
 from splunklib.modularinput.event_writer import EventWriter
-from io import BytesIO, TextIOWrapper
-
-try:
-    from splunklib.six.moves import cStringIO as StringIO
-except ImportError:
-    from splunklib.six import StringIO
-
-
-class EventTestCase(unittest.TestCase):
-    def test_event_without_enough_fields_fails(self):
-        """Check that events without data throw an error"""
-        with self.assertRaises(ValueError):
-            event = Event()
-            stream = StringIO()
-            event.write_to(stream)
-        self.assertTrue(True)
-
-    def test_xml_of_event_with_minimal_configuration(self):
-        """Generate XML from an event object with a small number of fields,
-        and see if it matches what we expect."""
-        stream = BytesIO()
-
-        event = Event(
-            data="This is a test of the emergency broadcast system.",
-            stanza="fubar",
-            time="%.3f" % 1372187084.000
-        )
-
-        event.write_to(stream)
-
-        constructed = ET.fromstring(stream.getvalue())
-        expected = ET.parse(data_open("data/event_minimal.xml")).getroot()
-
-        self.assertTrue(xml_compare(expected, constructed))
-
-    def test_xml_of_event_with_more_configuration(self):
-        """Generate XML from an even object with all fields set, see if
-        it matches what we expect"""
-        stream = BytesIO()
-
-        event = Event(
-            data="This is a test of the emergency broadcast system.",
-            stanza="fubar",
-            time="%.3f" % 1372274622.493,
-            host="localhost",
-            index="main",
-            source="hilda",
-            sourcetype="misc",
-            done=True,
-            unbroken=True
-        )
-        event.write_to(stream)
-
-        constructed = ET.fromstring(stream.getvalue())
-        expected = ET.parse(data_open("data/event_maximal.xml")).getroot()
-
-        self.assertTrue(xml_compare(expected, constructed))
-
-    def test_writing_events_on_event_writer(self):
-        """Write a pair of events with an EventWriter, and ensure that they
-        are being encoded immediately and correctly onto the output stream"""
-        out = BytesIO()
-        err = BytesIO()
-
-        ew = EventWriter(out, err)
-
-        e = Event(
-            data="This is a test of the emergency broadcast system.",
-            stanza="fubar",
-            time="%.3f" % 1372275124.466,
-            host="localhost",
-            index="main",
-            source="hilda",
-            sourcetype="misc",
-            done=True,
-            unbroken=True
-        )
-        ew.write_event(e)
-
-        found = ET.fromstring("%s</stream>" % out.getvalue().decode('utf-8'))
-        expected = ET.parse(data_open("data/stream_with_one_event.xml")).getroot()
 
-        self.assertTrue(xml_compare(expected, found))
-        self.assertEqual(err.getvalue(), b"")
 
+def test_event_without_enough_fields_fails(capsys):
+    """Check that events without data throw an error"""
+    with pytest.raises(ValueError), capsys.disabled():
+        event = Event()
+        event.write_to(sys.stdout)
+
+def test_xml_of_event_with_minimal_configuration(capsys):
+    """Generate XML from an event object with a small number of fields,
+    and see if it matches what we expect."""
+
+    event = Event(
+        data="This is a test of the emergency broadcast system.",
+        stanza="fubar",
+        time="%.3f" % 1372187084.000
+    )
+
+    event.write_to(sys.stdout)
+
+    captured = capsys.readouterr()
+    constructed = ET.fromstring(captured.out)
+    with data_open("data/event_minimal.xml") as data:
+        expected = ET.parse(data).getroot()
+
+        assert xml_compare(expected, constructed)
+
+def test_xml_of_event_with_more_configuration(capsys):
+    """Generate XML from an even object with all fields set, see if
+    it matches what we expect"""
+
+    event = Event(
+        data="This is a test of the emergency broadcast system.",
+        stanza="fubar",
+        time="%.3f" % 1372274622.493,
+        host="localhost",
+        index="main",
+        source="hilda",
+        sourcetype="misc",
+        done=True,
+        unbroken=True
+    )
+    event.write_to(sys.stdout)
+
+    captured = capsys.readouterr()
+
+    constructed = ET.fromstring(captured.out)
+    with data_open("data/event_maximal.xml") as data:
+        expected = ET.parse(data).getroot()
+
+        assert xml_compare(expected, constructed)
+
+def test_writing_events_on_event_writer(capsys):
+    """Write a pair of events with an EventWriter, and ensure that they
+    are being encoded immediately and correctly onto the output stream"""
+
+    ew = EventWriter(sys.stdout, sys.stderr)
+
+    e = Event(
+        data="This is a test of the emergency broadcast system.",
+        stanza="fubar",
+        time="%.3f" % 1372275124.466,
+        host="localhost",
+        index="main",
+        source="hilda",
+        sourcetype="misc",
+        done=True,
+        unbroken=True
+    )
+    ew.write_event(e)
+
+    captured = capsys.readouterr()
+
+    first_out_part = captured.out
+
+    with data_open("data/stream_with_one_event.xml") as data:
+        found = ET.fromstring("%s</stream>" % first_out_part)
+        expected = ET.parse(data).getroot()
+
+        assert xml_compare(expected, found)
+        assert captured.err == ""
+
+    ew.write_event(e)
+    ew.close()
+
+    captured = capsys.readouterr()
+    with data_open("data/stream_with_two_events.xml") as data:
+        found = ET.fromstring(first_out_part + captured.out)
+        expected = ET.parse(data).getroot()
+
+        assert xml_compare(expected, found)
+
+def test_error_in_event_writer():
+    """An event which cannot write itself onto an output stream
+    (such as because it doesn't have a data field set)
+    should write an error. Check that it does so."""
+
+    ew = EventWriter(sys.stdout, sys.stderr)
+    e = Event()
+    with pytest.raises(ValueError) as excinfo:
         ew.write_event(e)
-        ew.close()
+    assert str(excinfo.value) == "Events must have at least the data field set to be written to XML."
 
-        found = ET.fromstring(out.getvalue())
-        expected = ET.parse(data_open("data/stream_with_two_events.xml")).getroot()
+def test_logging_errors_with_event_writer(capsys):
+    """Check that the log method on EventWriter produces the
+    expected error message."""
 
-        self.assertTrue(xml_compare(expected, found))
+    ew = EventWriter(sys.stdout, sys.stderr)
 
-    def test_error_in_event_writer(self):
-        """An event which cannot write itself onto an output stream
-        (such as because it doesn't have a data field set)
-        should write an error. Check that it does so."""
-        out = BytesIO()
-        err = BytesIO()
+    ew.log(EventWriter.ERROR, "Something happened!")
 
-        ew = EventWriter(out, err)
-        e = Event()
-        try:
-            ew.write_event(e)
-            self.assertTrue(False)
-        except ValueError as e:
-            self.assertEqual("Events must have at least the data field set to be written to XML.", str(e))
+    captured = capsys.readouterr()
+    assert captured.err == "ERROR Something happened!\n"
 
-    def test_logging_errors_with_event_writer(self):
-        """Check that the log method on EventWriter produces the
-        expected error message."""
-        out = BytesIO()
-        err = BytesIO()
+def test_write_xml_is_sane(capsys):
+    """Check that EventWriter.write_xml_document writes sensible
+    XML to the output stream."""
 
-        ew = EventWriter(out, err)
+    ew = EventWriter(sys.stdout, sys.stderr)
 
-        ew.log(EventWriter.ERROR, "Something happened!")
-
-        self.assertEqual(b"ERROR Something happened!\n", err.getvalue())
-
-    def test_write_xml_is_sane(self):
-        """Check that EventWriter.write_xml_document writes sensible
-        XML to the output stream."""
-        out = BytesIO()
-        err = BytesIO()
-
-        ew = EventWriter(out, err)
-
-        expected_xml = ET.parse(data_open("data/event_maximal.xml")).getroot()
+    with data_open("data/event_maximal.xml") as data:
+        expected_xml = ET.parse(data).getroot()
 
         ew.write_xml_document(expected_xml)
-        found_xml = ET.fromstring(out.getvalue())
-
-        self.assertTrue(xml_compare(expected_xml, found_xml))
+        captured = capsys.readouterr()
+        found_xml = ET.fromstring(captured.out)
 
-if __name__ == "__main__":
-    unittest.main()
+        assert xml_compare(expected_xml, found_xml)
diff --git a/tests/modularinput/test_script.py b/tests/modularinput/test_script.py
