From 109a0d8da5666fb5adc07b03bc975747646992c0 Mon Sep 17 00:00:00 2001
From: Mike McLean <mikem@redhat.com>
Date: Nov 28 2018 20:26:29 +0000
Subject: PR#981: cli: add a param in watch_tasks to override KeyboardInterrupt output


Merges #981
https://pagure.io/koji/pull-request/981

Fixes: #976
https://pagure.io/koji/issue/976
Make output from watch_tasks customizable when ctrl+c is pressed

---

diff --git a/cli/koji_cli/lib.py b/cli/koji_cli/lib.py
index 162d6a1..1f677d2 100644
--- a/cli/koji_cli/lib.py
+++ b/cli/koji_cli/lib.py
@@ -276,11 +276,20 @@ def display_task_results(tasks):
             print('%s has not completed' % task_label)
 
 
-def watch_tasks(session, tasklist, quiet=False, poll_interval=60):
+def watch_tasks(session, tasklist, quiet=False, poll_interval=60, ki_handler=None):
     if not tasklist:
         return
     if not quiet:
         print("Watching tasks (this may be safely interrupted)...")
+    if ki_handler is None:
+        def ki_handler(progname, tasks, quiet):
+            if not quiet:
+                tlist = ['%s: %s' % (t.str(), t.display_state(t.info))
+                         for t in tasks.values() if not t.is_done()]
+                print(
+"""Tasks still running. You can continue to watch with the '%s watch-task' command.
+Running Tasks:
+%s""" % (progname, '\n'.join(tlist)))
     sys.stdout.flush()
     rv = 0
     try:
@@ -317,14 +326,9 @@ def watch_tasks(session, tasklist, quiet=False, poll_interval=60):
             sys.stdout.flush()
             time.sleep(poll_interval)
     except KeyboardInterrupt:
-        if tasks and not quiet:
+        if tasks:
             progname = os.path.basename(sys.argv[0]) or 'koji'
-            tlist = ['%s: %s' % (t.str(), t.display_state(t.info))
-                            for t in tasks.values() if not t.is_done()]
-            print( \
-"""Tasks still running. You can continue to watch with the '%s watch-task' command.
-Running Tasks:
-%s""" % (progname, '\n'.join(tlist)))
+            ki_handler(progname, tasks, quiet)
         raise
     return rv
 
diff --git a/tests/test_cli/test_watch_tasks.py b/tests/test_cli/test_watch_tasks.py
index fafcc4d..e95407b 100644
--- a/tests/test_cli/test_watch_tasks.py
+++ b/tests/test_cli/test_watch_tasks.py
@@ -121,6 +121,32 @@ Running Tasks:
 ''' % (os.path.basename(sys.argv[0]) or 'koji'))
         self.assertMultiLineEqual(stdout.getvalue(), expected)
 
+    @mock.patch('time.sleep')
+    @mock.patch('sys.stdout', new_callable=six.StringIO)
+    def test_watch_tasks_with_keyboardinterrupt_handler(self, stdout, sleep):
+        """Raise KeyboardInterrupt inner watch_tasks with a ki_handler"""
+        cfile = os.path.dirname(__file__) + '/data/calls/watchtasks2.json'
+        with open(cfile) as fp:
+            cdata = json.load(fp)
+        self.session.load_calls(cdata)
+        sleep.side_effect = [None] * 10 + [KeyboardInterrupt]
+
+        def customized_handler(progname, tasks, quiet):
+            print('some output')
+
+        with self.assertRaises(KeyboardInterrupt):
+            # watch_tasks catches and re-raises it to display a message
+            watch_tasks(self.session, [1208], quiet=False, poll_interval=5,
+                        ki_handler=customized_handler)
+        expected = ('''Watching tasks (this may be safely interrupted)...
+1208 build (f24, /users/mikem/fake.git:master): free
+1208 build (f24, /users/mikem/fake.git:master): free -> open (builder-01)
+  1209 buildSRPMFromSCM (/users/mikem/fake.git:master): free
+  1209 buildSRPMFromSCM (/users/mikem/fake.git:master): free -> open (builder-01)
+some output
+''')
+        self.assertMultiLineEqual(stdout.getvalue(), expected)
+
 
 if __name__ == '__main__':
     unittest.main()