Cleaner naming and semantics for custom search command input streams

Shindora · Sep 4, 2020 · 8dec167 · 8dec167
1 parent 253784c
commit 8dec167
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 14 deletions.
diff --git a/splunklib/searchcommands/generating_command.py b/splunklib/searchcommands/generating_command.py
@@ -204,7 +204,7 @@ def _execute(self, ifile, process):
 
         """
         if self._protocol_version == 2:
-            result = self._read_chunk(ifile)
+            result = self._read_chunk(self._as_binary_stream(ifile))
 
             if not result:
                 return

diff --git a/splunklib/searchcommands/search_command.py b/splunklib/searchcommands/search_command.py
@@ -656,7 +656,7 @@ def _process_protocol_v2(self, argv, ifile, ofile):
         # noinspection PyBroadException
         try:
             debug('Reading metadata')
-            metadata, body = self._read_chunk(ifile)
+            metadata, body = self._read_chunk(self._as_binary_stream(ifile))
 
             action = getattr(metadata, 'action', None)
 
@@ -850,15 +850,20 @@ def _execute(self, ifile, process):
         self.finish()
 
     @staticmethod
-    def _read_chunk(ifile):
-        # noinspection PyBroadException
+    def _as_binary_stream(ifile):
         if six.PY2:
-            istream = ifile
-        else:
-            try:
-                istream = ifile.buffer
-            except AttributeError as error:
-                raise RuntimeError('Failed to get underlying buffer: {}'.format(error))
+            return ifile
+
+        try:
+            return ifile.buffer
+        except AttributeError as error:
+            raise RuntimeError('Failed to get underlying buffer: {}'.format(error))
+
+    @staticmethod
+    def _read_chunk(istream):
+        # noinspection PyBroadException
+        if not six.PY2:
+            assert issubclass(type(istream), (io.BufferedIOBase, io.RawIOBase)), 'Stream must be binary'
 
         try:
             header = istream.readline()
@@ -930,9 +935,10 @@ def _records_protocol_v1(self, ifile):
             yield record
 
     def _records_protocol_v2(self, ifile):
+        istream = self._as_binary_stream(ifile)
 
         while True:
-            result = self._read_chunk(ifile)
+            result = self._read_chunk(istream)
 
             if not result:
                 return

diff --git a/tests/searchcommands/test_search_command.py b/tests/searchcommands/test_search_command.py
@@ -38,10 +38,10 @@
 import pytest
 
 def build_command_input(getinfo_metadata, execute_metadata, execute_body):
-    header = ('chunked 1.0,{},0\n{}'.format(len(six.ensure_binary(getinfo_metadata)), getinfo_metadata) +
-              'chunked 1.0,{},{}\n{}{}'.format(len(six.ensure_binary(execute_metadata)), len(six.ensure_binary(execute_body)), execute_metadata, execute_body))
+    input = ('chunked 1.0,{},0\n{}'.format(len(six.ensure_binary(getinfo_metadata)), getinfo_metadata) +
+             'chunked 1.0,{},{}\n{}{}'.format(len(six.ensure_binary(execute_metadata)), len(six.ensure_binary(execute_body)), execute_metadata, execute_body))
 
-    ifile = BytesIO(six.ensure_binary(header))
+    ifile = BytesIO(six.ensure_binary(input))
 
     if not six.PY2:
         ifile = TextIOWrapper(ifile)