script mode testing and fixes

tests/test_executenode.py

@@ -156,17 +156,28 @@ class TestExecuteNode(unittest2.TestCase):
         self.assertEqual(nodeb.run(cmd=["pwd", ExecuteNode.PIPE, "cat"], cwd="/tmp/space test"), ["/tmp/space test"])
         self.assertEqual(nodeb.run(cmd=["cat", ExecuteNode.PIPE, "pwd"], cwd="/tmp/space test"), ["/tmp/space test"])
 
-    # #
-    # def test_script(self):
-    #
-    #     def stdout_handler(line):
-    #         print("handle: " + line)
-    #
-    #     nodea=ExecuteNode(debug_output=True, ssh_to="localhost")
-    #
-    #     cmd_pipe=nodea.script(lines=["echo line1", "echo line 2"], stdout_handler=stdout_handler)
-    #     cmd_pipe.execute()
+    def test_script_handlers(self):
 
+        results=[]
+        nodea=ExecuteNode(debug_output=True)
+        cmd_pipe=nodea.script(lines=["echo line1", "echo line2 1>&2", "exit 123"],
+                              stdout_handler=lambda line: results.append(line),
+                              stderr_handler=lambda line: results.append(line),
+                              exit_handler=lambda exit_code: results.append(exit_code),
+                              valid_exitcodes=[123]
+                              )
+        cmd_pipe.execute()
+
+        self.assertEqual(results, ["line1", "line2", 123 ])
+
+    def test_script_defaults(self):
+
+        def handler(line):
+            pass
+
+        nodea=ExecuteNode(debug_output=True, ssh_to="localhost")
+        cmd_pipe=nodea.script(lines=["echo test"], stdout_handler=handler)
+        cmd_pipe.execute()
 
 
 

zfs_autobackup/CmdPipe.py

@@ -143,13 +143,6 @@ class CmdPipe:
             # read line and call appropriate handlers
 
             for item in self.items:
-                if item.process.stderr in read_ready:
-                    line = item.process.stderr.readline().decode('utf-8').rstrip()
-                    if line != "":
-                        item.stderr_handler(line)
-                    else:
-                        eof_count = eof_count + 1
-
                 if item.process.stdout in read_ready:
                     line = item.process.stdout.readline().decode('utf-8').rstrip()
                     if line != "":
@@ -159,6 +152,14 @@ class CmdPipe:
                         if item.next:
                             item.next.process.stdin.close()
 
+                if item.process.stderr in read_ready:
+                    line = item.process.stderr.readline().decode('utf-8').rstrip()
+                    if line != "":
+                        item.stderr_handler(line)
+                    else:
+                        eof_count = eof_count + 1
+
+
                 if item.process.poll() is not None:
                     done_count = done_count + 1
 

zfs_autobackup/ExecuteNode.py

@@ -183,7 +183,7 @@ class ExecuteNode(LogStub):
         else:
             return output_lines
 
-    def script(self, lines, inp=None, valid_exitcodes=None, readonly=False, hide_errors=False):
+    def script(self, lines, inp=None, stdout_handler=None, stderr_handler=None, exit_handler=None, valid_exitcodes=None, readonly=False, hide_errors=False):
         """Run a multiline script on the node.
 
         This is much more low level than run() and allows for finer grained control.
@@ -205,7 +205,6 @@ class ExecuteNode(LogStub):
 
         """
 
-
         # create new pipe?
         if not isinstance(inp, CmdPipe):
             cmd_pipe = CmdPipe(self.readonly, inp)
@@ -213,20 +212,31 @@ class ExecuteNode(LogStub):
             # add stuff to existing pipe
             cmd_pipe = inp
 
+        internal_stdout_handler=None
+        if stdout_handler is not None:
+            if self.debug_output:
+                def internal_stdout_handler(line):
+                    self.debug("STDOUT > " + line.rstrip())
+                    stdout_handler(line)
 
-        def stderr_handler(line):
+        def internal_stderr_handler(line):
             self._parse_stderr(line, hide_errors)
+            if stderr_handler is not None:
+                stderr_handler(line)
 
         # exit code hanlder
         if valid_exitcodes is None:
             valid_exitcodes = [0]
 
-        def exit_handler(exit_code):
+        def internal_exit_handler(exit_code):
             if self.debug_output:
                 self.debug("EXIT   > {}".format(exit_code))
 
+            if exit_handler is not None:
+                exit_handler(exit_code)
+
             if (valid_exitcodes != []) and (exit_code not in valid_exitcodes):
-                self.error("Script returned exit code {} (valid codes: {})".format(cmd_item, exit_code, valid_exitcodes))
+                self.error("Script returned exit code {} (valid codes: {})".format(exit_code, valid_exitcodes))
                 return False
 
             return True
@@ -248,7 +258,7 @@ class ExecuteNode(LogStub):
         cmd.append("\n".join(lines))
 
         # add shell command and handlers to pipe
-        cmd_item=CmdItem(cmd=cmd, readonly=readonly, stderr_handler=stderr_handler, exit_handler=exit_handler, shell=self.is_local())
+        cmd_item=CmdItem(cmd=cmd, readonly=readonly, stderr_handler=internal_stderr_handler, exit_handler=internal_exit_handler, stdout_handler=internal_stdout_handler, shell=self.is_local())
         cmd_pipe.add(cmd_item)
 
         self.debug("SCRIPT > {}".format(cmd_pipe))