From 511dfc8fddf2280e2fcb61853c4df6791294f157 Mon Sep 17 00:00:00 2001
From: Yuming Zhu <yzhu@redhat.com>
Date: Tue, 4 Sep 2018 17:39:46 +0800
Subject: [PATCH] hub: [getPackageID] add strict behavior

---
 hub/kojihub.py                      | 18 +++++++----
 tests/test_hub/test_getPackageID.py | 47 +++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 6 deletions(-)
 create mode 100644 tests/test_hub/test_getPackageID.py

diff --git a/hub/kojihub.py b/hub/kojihub.py
index b4071034..789f6b31 100644
--- a/hub/kojihub.py
+++ b/hub/kojihub.py
@@ -10221,14 +10221,20 @@ class RootExports(object):
     getTagID = staticmethod(get_tag_id)
     getTag = staticmethod(get_tag)
 
-    def getPackageID(self, name):
-        c = context.cnx.cursor()
-        q = """SELECT id FROM package WHERE name=%(name)s"""
-        c.execute(q, locals())
-        r = c.fetchone()
+    def getPackageID(self, name, strict=False):
+        """Get package ID by name.
+        If package doesn't exist, return None, unless strict is True in which
+        case an exception is raised."""
+        query = QueryProcessor(tables=['package'],
+                               columns=['id'],
+                               clauses=['name=%(name)s'],
+                               values=locals())
+        r = query.executeOne()
         if not r:
+            if strict:
+                raise koji.GenericError('Invalid package name: %s' % name)
             return None
-        return r[0]
+        return r['id']
 
     getPackage = staticmethod(lookup_package)
 
diff --git a/tests/test_hub/test_getPackageID.py b/tests/test_hub/test_getPackageID.py
new file mode 100644
index 00000000..8172bbc3
--- /dev/null
+++ b/tests/test_hub/test_getPackageID.py
@@ -0,0 +1,47 @@
+from __future__ import absolute_import
+
+import mock
+
+from .utils import DBQueryTestCase
+
+import koji
+import kojihub
+
+
+class TestGetPackageID(DBQueryTestCase):
+    maxDiff = None
+
+    def test_getPackageID(self):
+        self.qp_execute_return_value = [{'id': 1}]
+        rv = kojihub.RootExports().getPackageID('koji')
+        self.assertEqual(len(self.queries), 1)
+        self.assertLastQueryEqual(tables=['package'],
+                                  columns=['id'],
+                                  clauses=['name=%(name)s'],
+                                  values={'name': 'koji',
+                                          'strict': False,
+                                          'self': mock.ANY})
+        self.assertEqual(rv, 1)
+
+    def test_getPackageID_strict(self):
+        with self.assertRaises(koji.GenericError) as cm:
+            kojihub.RootExports().getPackageID('invalidpkg', strict=True)
+        self.assertLastQueryEqual(tables=['package'],
+                                  columns=['id'],
+                                  clauses=['name=%(name)s'],
+                                  values={'name': 'invalidpkg',
+                                          'strict': True,
+                                          'self': mock.ANY})
+        self.assertEqual(cm.exception.args[0],
+                         'Invalid package name: invalidpkg')
+
+    def test_getPackageID_None(self):
+        rv = kojihub.RootExports().getPackageID('invalidpkg')
+        self.assertEqual(len(self.queries), 1)
+        self.assertLastQueryEqual(tables=['package'],
+                                  columns=['id'],
+                                  clauses=['name=%(name)s'],
+                                  values={'name': 'invalidpkg',
+                                          'strict': False,
+                                          'self': mock.ANY})
+        self.assertIsNone(rv)