From 30ed89e485147997161e4029a09f1c369e6d93a2 Mon Sep 17 00:00:00 2001
From: Yuming Zhu <yzhu@redhat.com>
Date: Nov 28 2018 20:25:57 +0000
Subject: cli: add a param in watch_tasks to override KeyboardInterrupt output


---

diff --git a/cli/koji_cli/lib.py b/cli/koji_cli/lib.py
index 162d6a1..3d4f2a2 100644
--- a/cli/koji_cli/lib.py
+++ b/cli/koji_cli/lib.py
@@ -276,7 +276,7 @@ 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:
@@ -319,12 +319,15 @@ def watch_tasks(session, tasklist, quiet=False, poll_interval=60):
     except KeyboardInterrupt:
         if tasks and not quiet:
             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( \
+            if ki_handler is None:
+                def ki_handler(progname, tasks):
+                    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)
         raise
     return rv
 
diff --git a/tests/test_cli/test_watch_tasks.py b/tests/test_cli/test_watch_tasks.py
index fafcc4d..c5bb44e 100644
--- a/tests/test_cli/test_watch_tasks.py
+++ b/tests/test_cli/test_watch_tasks.py
@@ -121,6 +121,31 @@ 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):
+            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()