apply get_header_field enhancements from kobo and update unit test

koji/__init__.py

@@ -876,19 +876,44 @@ def get_rpm_header(f, ts=None):
         fo.close()
     return hdr
 
-def get_header_field(hdr, name):
+
+def get_header_field(hdr, name, src_arch=False):
     """Extract named field from an rpm header"""
     if not RPM_SUPPORTS_OPTIONAL_DEPS and name in ('SUGGESTNAME', 'SUGGESTVERSION', 'SUGGESTFLAGS',
                                                    'ENHANCENAME', 'ENHANCEVERSION', 'ENHANCEFLAGS',
                                                    'SUPPLEMENTNAME', 'SUPPLEMENTVERSION', 'SUPPLEMENTFLAGS',
                                                    'RECOMMENDNAME', 'RECOMMENDVERSION', 'RECOMMENDFLAGS'):
         return []
-    idx = getattr(rpm, "RPMTAG_%s" % name.upper(), None)
-    if idx is None:
-        raise GenericError("No such rpm header field: %s" % name)
-    return hdr[idx]
 
-def get_header_fields(X, fields):
+    if (src_arch and name.lower() == "arch"
+                and get_header_field(hdr, "sourcepackage")):
+        # return "src" or "nosrc" arch instead of build arch for src packages
+        if (get_header_field(hdr, "nosource")
+                    or get_header_field(hdr, "nopatch")):
+            return "nosrc"
+        return "src"
+
+    hdr_key = getattr(rpm, "RPMTAG_%s" % name.upper(), None)
+    if hdr_key is None:
+        # HACK: nosource and nopatch may not be in exported rpm tags
+        if name == "nosource":
+            hdr_key = 1051
+        elif name == "nopatch":
+            hdr_key = 1052
+        else:
+            raise GenericError("No such rpm header field: %s" % name)
+
+    result = hdr[hdr_key]
+    if name in ("nosource", "nopatch"):
+        # HACK: workaround for https://bugzilla.redhat.com/show_bug.cgi?id=991329
+        if result is None:
+            result = []
+        elif isinstance(result, (int, long)):
+            result = [result]
+    return result
+
+
+def get_header_fields(X, fields, src_arch=False):
     """Extract named fields from an rpm header and return as a dictionary
 
     X may be either the rpm header or the rpm filename
@@ -899,7 +924,7 @@ def get_header_fields(X, fields):
         hdr = X
     ret = {}
     for f in fields:
-        ret[f] = get_header_field(hdr, f)
+        ret[f] = get_header_field(hdr, f, src_arch=src_arch)
     return ret
 
 def parse_NVR(nvr):

tests/data/specs/test-nopatch._spec

@@ -0,0 +1,13 @@
+Name:           test-nopatch
+Version:        1
+Release:        1%{?dist}
+Summary:        Testing source arch header fields
+
+License:        none
+Patch0:         secret.patch
+Nopatch:        0
+
+%description
+...
+
+%files

tests/data/specs/test-nosrc._spec

@@ -0,0 +1,13 @@
+Name:           test-nosrc
+Version:        1
+Release:        1%{?dist}
+Summary:        Testing source arch header fields
+
+License:        none
+Source0:        secret.key
+Nosource:       0
+
+%description
+...
+
+%files

tests/data/specs/test-src._spec

@@ -0,0 +1,11 @@
+Name:		test-src
+Version:	1
+Release:	1%{?dist}
+Summary:	Testing source arch header fields
+
+License:        none
+
+%description
+...
+
+%files

tests/test_parsers.py

@@ -113,6 +113,7 @@ class INITTestCase(unittest.TestCase):
 
 class HeaderTestCase(unittest.TestCase):
     rpm_path = os.path.join(os.path.dirname(__file__), 'data/rpms/test-deps-1-1.fc24.x86_64.rpm')
+    rpmdir = os.path.join(os.path.dirname(__file__), 'data/rpms')
 
     def setUp(self):
         self.fd = open(self.rpm_path)
@@ -141,5 +142,33 @@ class HeaderTestCase(unittest.TestCase):
         self.assertEqual(['REQUIRES'], koji.get_header_fields(hdr, ['REQUIRES']).keys())
 
 
+    def test_get_header_field_src(self):
+        srpm = os.path.join(self.rpmdir, 'test-src-1-1.fc24.src.rpm')
+
+        # without src_arch, should return the build arch (x86_64)
+        data =  koji.get_header_fields(srpm, ['arch'])
+        self.assertEqual(data['arch'], 'x86_64')
+
+        # eith src_arch, should return src
+        data =  koji.get_header_fields(srpm, ['arch'], src_arch=True)
+        self.assertEqual(data['arch'], 'src')
+
+
+    def test_get_header_field_nosrc(self):
+        srpm1 = os.path.join(self.rpmdir, 'test-nosrc-1-1.fc24.nosrc.rpm')
+        srpm2 = os.path.join(self.rpmdir, 'test-nopatch-1-1.fc24.nosrc.rpm')
+
+        # without src_arch, should return the build arch (x86_64)
+        for srpm in srpm1, srpm2:
+            data =  koji.get_header_fields(srpm, ['arch'])
+            self.assertEqual(data['arch'], 'x86_64')
+
+        # eith src_arch, should return nosrc
+        for srpm in srpm1, srpm2:
+            data =  koji.get_header_fields(srpm, ['arch'], src_arch=True)
+            self.assertEqual(data['arch'], 'nosrc')
+
+
+
 if __name__ == '__main__':
     unittest.main()