Generatorfix (#181)

* fix dvpl-7207

* add condition change only for v2

* fix unittest

* fix generator unittest

* remove comment out code

* small fix: make \n check optional

* small format fix

Author: ljiang1

splunklib/searchcommands/generating_command.py

@@ -193,6 +193,18 @@ def _execute(self, ifile, process):
         :return: `None`.
 
         """
+        if self._protocol_version == 2:
+            result = self._read_chunk(ifile)
+
+            if not result:
+                return
+
+            metadata, body = result
+            action = getattr(metadata, 'action', None)
+
+            if action != 'execute':
+                raise RuntimeError('Expected execute action, not {}'.format(action))
+
         self._record_writer.write_records(self.generate())
         self.finish()
 

splunklib/searchcommands/search_command.py

@@ -19,6 +19,7 @@
 # Absolute imports
 
 from collections import namedtuple
+
 try:
     from collections import OrderedDict  # must be python 2.7
 except ImportError:
@@ -27,6 +28,7 @@
 from cStringIO import StringIO
 from itertools import chain, ifilter, imap, islice, izip
 from logging import _levelNames, getLevelName, getLogger
+
 try:
     from shutil import make_archive
 except ImportError:
@@ -47,7 +49,7 @@
 
 # Relative imports
 
-from . internals import (
+from .internals import (
     CommandLineParser,
     CsvDialect,
     InputHeader,
@@ -63,6 +65,7 @@
 from . import Boolean, Option, environment
 from ..client import Service
 
+
 # ----------------------------------------------------------------------------------------------------------------------
 
 # P1 [ ] TODO: Log these issues against ChunkedExternProcessor
@@ -89,6 +92,7 @@ class SearchCommand(object):
     """ Represents a custom search command.
 
     """
+
     def __init__(self):
 
         # Variables that may be used, but not altered by derived classes
@@ -873,8 +877,10 @@ def _read_chunk(ifile):
         # if body_length <= 0:
         #     return metadata, ''
 
+        body = ""
         try:
-            body = ifile.read(body_length)
+            if body_length > 0:
+                body = ifile.read(body_length)
         except Exception as error:
             raise RuntimeError('Failed to read body of length {}: {}'.format(body_length, error))