libpipe 0.2__tar.gz → 0.4__tar.gz

{libpipe-0.2 → libpipe-0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: libpipe
-Version: 0.2
+Version: 0.4
 Summary: 
 Home-page: https://gitlab.com/flomertens/libpipe
 License: GPL-3.0-or-later

{libpipe-0.2 → libpipe-0.4}/libpipe/tests/test_worker.py

@@ -1,6 +1,7 @@
 import os
 import string
 import random
+import functools
 
 from unittest import mock
 
@@ -173,6 +174,37 @@ def test_done_callback():
     assert done
 
 
+def test_task_done_callback():
+    w = worker.WorkerPool([worker.localhost_shortname], 'Test', max_time=2)
+
+    done = 0
+
+    def callback(task):
+        print(task)
+        nonlocal done
+        assert isinstance(task, worker.Task)
+        if task.name == 'Error':
+            assert task.returncode == 1
+        else:    
+            assert task.returncode == 0
+        done += 1
+
+    def callback_with_extra(name, task):
+        assert name == 'test'
+        callback(task)
+
+    w.add('echo 0', task_done_callback=callback, name='Success1')
+    w.add('echo 0', task_done_callback=callback, name='Success2')
+    w.add('echo 3; cat non_existing_file', task_done_callback=functools.partial(callback_with_extra, 'test'), name='Error')
+    w.add('echo 4;', task_done_callback=functools.partial(callback_with_extra, 'test_err'), name='CBK Error')
+
+    success, errors = w.execute()
+
+    assert len(success) == 4
+    assert len(errors) == 0
+    assert done == 3
+
+
 class FakeSSHClient(worker.LocalClient):
 
     def __init__(self, host, user=None, force_sync=True, password=None):
@@ -229,3 +261,9 @@ def test_line_too_long(tmp_path):
     assert has_done
     assert len(errors) == 0
     assert len(success) == 3
+
+
+def test_client():
+    client = worker.Client('localhost')
+
+    task = worker.Task('Check size', 'du -hs *')

{libpipe-0.2 → libpipe-0.4}/libpipe/worker.py

@@ -91,7 +91,7 @@ class Task(object):
         returncode (int): Return code. None if task was not executed.
     """
 
-    def __init__(self, name, command, output_file=None, done_callback=None, 
+    def __init__(self, name, command, output_file=None, done_callback=None, task_done_callback=None,
                  run_on_host=None, n_try=3, pid_filename=None, keep_pid_file_on_error=False):
         self.name = name
         self.command = command
@@ -99,7 +99,8 @@ class Task(object):
         self.fd = None
         self.returncode = None
         self.n_try = 0
-        self.done_callback = done_callback
+        self._done_callback = done_callback
+        self._task_done_callback = task_done_callback
         self.run_on_host = run_on_host
         self.status = TaskStatus.WAITING
         self.process = None
@@ -108,6 +109,21 @@ class Task(object):
         self.pid_filename = pid_filename
         self.keep_pid_file_on_error = keep_pid_file_on_error
 
+    def done_callback(self):
+        err_msg = ''
+        if self._done_callback is not None:
+            try:
+                self._done_callback()
+            except Exception as exc:
+                err_msg = f'Error executing return callback of task {self.name}: {str(exc)}\n'
+        if self._task_done_callback is not None:
+            try:
+                self._task_done_callback(self)
+            except Exception as exc:
+                err_msg = f'Error executing return callback of task {self.name}: {str(exc)}\n'
+
+        return err_msg
+
     def init_log(self):
         if self.output_file is not None:
             self.fd = open(self.output_file, 'w')
@@ -232,7 +248,8 @@ class Client(object):
             print(f'Starting client {self.host} ...')
             try:
                 self.conn, self.client = await asyncssh.create_connection(SSHClient, self.host, username=self.user,
-                                                                          password=self.password)
+                                                                          password=self.password, agent_forwarding=True)
+                self.conn.set_keepalive(60, 3)
             except Exception:
                 raise
             finally:
@@ -342,11 +359,7 @@ class Worker(object):
             self.status = WorkerStatus.IDLE
             return False, err_msg
 
-        if task.done_callback is not None:
-            try:
-                task.done_callback()
-            except Exception as exc:
-                err_msg = f'Error executing return callback of task {self.name}: {str(exc)}\n'
+        err_msg = task.done_callback()
 
         self.status = WorkerStatus.IDLE
 
@@ -421,7 +434,7 @@ class WorkerPool(object):
                     return worker
             await asyncio.sleep(0.1)
 
-    def add(self, command, name=None, output_file=None, done_callback=None, 
+    def add(self, command, name=None, output_file=None, task_done_callback=None, done_callback=None, 
             run_on_host=None, n_try=3, pid_filename=None, keep_pid_file_on_error=False):
         """Add a command to execute by a worker in the pool. Optionally output result in output_file.
 
@@ -447,8 +460,8 @@ class WorkerPool(object):
             return
 
         self.tasks.append(Task(name, command, output_file=output_file, done_callback=done_callback,
-                               run_on_host=run_on_host, n_try=n_try, pid_filename=pid_filename, 
-                               keep_pid_file_on_error=keep_pid_file_on_error))
+                               task_done_callback=task_done_callback, run_on_host=run_on_host, n_try=n_try, 
+                               pid_filename=pid_filename, keep_pid_file_on_error=keep_pid_file_on_error))
 
     def get_all_tasks(self):
         tasks = collections.defaultdict(list)

{libpipe-0.2 → libpipe-0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "libpipe"
-version = "0.2"
+version = "0.4"
 description = ""
 authors = ["\"Florent Mertens\" <\"florent.mertens@gmail.com\">"]
 license = "GPL-3.0-or-later"

{libpipe-0.2 → libpipe-0.4}/setup.py

@@ -18,7 +18,7 @@ entry_points = \
 
 setup_kwargs = {
     'name': 'libpipe',
-    'version': '0.2',
+    'version': '0.4',
     'description': '',
     'long_description': 'Libpipe\n=======\n\nCommon pipeline framework and library\n\nInstallation\n------------\n\nlibpipe can be installed via pip:\n\n    $ pip install libpipe\n\nand requires Python 3.6.0 or higher.\n',
     'author': '"Florent Mertens"',