From 53d57f2a6f4d0acbbc626ff04d5aa5318c2721d2 Mon Sep 17 00:00:00 2001
From: Tomas Kopecek <tkopecek@redhat.com>
Date: Sep 23 2020 12:58:07 +0000
Subject: hub: extract cg_name from build in policy_get_cgs


Fixes: https://pagure.io/koji/issue/2251

---

diff --git a/hub/kojihub.py b/hub/kojihub.py
index 8537990..a722041 100644
--- a/hub/kojihub.py
+++ b/hub/kojihub.py
@@ -9277,6 +9277,18 @@ def policy_get_cgs(data):
         cgs = [lookup_name('content_generator', cg, strict=True)
                for cg in data['cg_list']]
         return set(cgs)
+
+    if 'build' in data:
+        if isinstance(data['build'], (str, int)):
+            build = get_build(data['build'], strict=True)
+        elif isinstance(data['build'], dict):
+            build = data['build']
+        else:
+            raise koji.GenericError('Unexpected policy_data["build"] type %s' %
+                                    type(data['build']))
+        if 'cg_name' in build:
+            return set([build['cg_name']])
+
     # otherwise try buildroot data
     cgs = set()
     for br_id in policy_get_brs(data):
diff --git a/tests/test_hub/test_policy_tests.py b/tests/test_hub/test_policy_tests.py
index bb153d7..7daab8e 100644
--- a/tests/test_hub/test_policy_tests.py
+++ b/tests/test_hub/test_policy_tests.py
@@ -128,7 +128,14 @@ class TestPolicyGetCGs(unittest.TestCase):
         self.assertEqual(result, expect)
         self.list_rpms.assert_called_once_with(buildID=42)
         self.list_archives.assert_called_once_with(buildID=42)
-        self.get_build.assert_called_once_with('NVR', strict=True)
+        self.get_build.assert_called_with('NVR', strict=True)
+
+    def test_policy_get_cg_from_build(self):
+        self.get_build.return_value = {'id': 42, 'cg_name': 'cgname'}
+        # let's see...
+        result = kojihub.policy_get_cgs({'build': 'NVR'})
+        self.assertEqual(result, set(['cgname']))
+        self.get_build.assert_called_with('NVR', strict=True)
 
     def test_policy_get_cg_from_cgs(self):
         data = {