From 910ed50456ea41995989d59a5a469799ba8d4e5b Mon Sep 17 00:00:00 2001
From: Ondrej Nosek <onosek@redhat.com>
Date: Wed, 15 May 2019 18:53:30 +0200
Subject: [PATCH] Custom handler for koji watch_tasks

Output text during rhpkg/fedpkg build process states that there
is a 'watch_task' subcommand. When 'koji_cli' library is imported
in rhpkg/fedpkg tool, it shows that command is named
'rhpkg/fedpkg watch_task' instead of 'brew/koji watch_task'. Custom
handler replaces the internal one inside koji_cli library.
Additional fix in rhpkg is needed after this change is released.

Relates: rhbz#1570921
Relates: COMPOSE-2809

Signed-off-by: Ondrej Nosek <onosek@redhat.com>
---
 pyrpkg/cli.py     |  6 +++++-
 pyrpkg/utils.py   | 21 +++++++++++++++++++++
 tests/test_cli.py | 39 ++++++++++++++++++++++-----------------
 3 files changed, 48 insertions(+), 18 deletions(-)

diff --git a/pyrpkg/cli.py b/pyrpkg/cli.py
index 7b83a44..72be1cf 100644
--- a/pyrpkg/cli.py
+++ b/pyrpkg/cli.py
@@ -1667,7 +1667,11 @@ see API KEY section of copr-cli(1) man page.
         if self.args.dry_run:
             self.log.info('DRY-RUN: Watch tasks: %s', task_ids)
         else:
-            return koji_cli.lib.watch_tasks(self.cmd.kojisession, task_ids)
+            return koji_cli.lib.watch_tasks(
+                self.cmd.kojisession,
+                task_ids,
+                ki_handler=utils.make_koji_watch_tasks_handler(self.cmd.build_client)
+            )
 
     def extract_greenwave_url(self):
         greenwave_url = None
diff --git a/pyrpkg/utils.py b/pyrpkg/utils.py
index 2268e6f..37a39e2 100644
--- a/pyrpkg/utils.py
+++ b/pyrpkg/utils.py
@@ -128,3 +128,24 @@ def validate_module_build_optional(optional_arg):
             'The "{0}" optional argument is reserved to built-in arguments'.format(key))
 
     return (key, value)
+
+
+def make_koji_watch_tasks_handler(progname):
+    def koji_watch_tasks_handler(_, tasks, quiet):
+        """
+        Displays information about running tasks and says how to watch them.
+        Unlike the default version at koji library it overrides progname
+        to show brew, koji or other build client.
+        """
+        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)))
+
+    # Save reference of the handler during first time use.
+    # It guarantees that the same object is always returned (it allows unittest to pass).
+    global handler_reference
+    if 'handler_reference' not in globals():
+        handler_reference = koji_watch_tasks_handler
+    return handler_reference
diff --git a/tests/test_cli.py b/tests/test_cli.py
index 785e103..b3e0718 100644
--- a/tests/test_cli.py
+++ b/tests/test_cli.py
@@ -3235,23 +3235,28 @@ class TestBuildPackage(FakeKojiCreds, CliTestCase):
         mock_build_api = None
 
         with patch('koji_cli.lib.watch_tasks') as watch_tasks:
-            with patch('sys.argv', new=cli_cmd):
-                cli = self.new_cli(cfg=config_file)
-                if sub_command == 'build':
-                    mock_build_api = session.build
-                    cli.build()
-                elif sub_command == 'scratch-build':
-                    mock_build_api = session.build
-                    cli.scratch_build()
-                elif sub_command == 'chain-build':
-                    mock_build_api = session.chainBuild
-                    cli.chainbuild()
-
-            if '--nowait' in cli_cmd:
-                watch_tasks.assert_not_called()
-            else:
-                watch_tasks.assert_called_once_with(
-                    session, [mock_build_api.return_value])
+            with patch('pyrpkg.utils.make_koji_watch_tasks_handler') as mock_ki:
+                with patch('sys.argv', new=cli_cmd):
+                    cli = self.new_cli(cfg=config_file)
+                    if sub_command == 'build':
+                        mock_build_api = session.build
+                        cli.build()
+                    elif sub_command == 'scratch-build':
+                        mock_build_api = session.build
+                        cli.scratch_build()
+                    elif sub_command == 'chain-build':
+                        mock_build_api = session.chainBuild
+                        cli.chainbuild()
+
+                if '--nowait' in cli_cmd:
+                    watch_tasks.assert_not_called()
+                else:
+                    watch_tasks.assert_called_once_with(
+                        session,
+                        [mock_build_api.return_value],
+                        ki_handler=mock_ki.return_value
+                    )
+                    self.assertEqual(mock_ki.call_args, (("koji",),))
 
         mock_build_api.assert_called_once()
 
-- 
2.20.1