From 8eec9fa56f0ad8991ec5581abbac4e65e2411f01 Mon Sep 17 00:00:00 2001
From: Scott McMillan <smcmillan@nvidia.com>
Date: Wed, 18 Sep 2024 16:52:40 -0500
Subject: [PATCH] Fix missing ubuntu24 distro override

---
 hpccm/primitives/baseimage.py |  2 ++
 test/test_baseimage.py        | 49 +++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+)

diff --git a/hpccm/primitives/baseimage.py b/hpccm/primitives/baseimage.py
index 6c1efec9..a714f996 100644
--- a/hpccm/primitives/baseimage.py
+++ b/hpccm/primitives/baseimage.py
@@ -121,6 +121,8 @@ def __init__(self, **kwargs):
             hpccm.config.set_linux_distro('ubuntu20')
         elif self.__distro == 'ubuntu22':
             hpccm.config.set_linux_distro('ubuntu22')
+        elif self.__distro == 'ubuntu24':
+            hpccm.config.set_linux_distro('ubuntu24')
         elif self.__distro == 'centos':
             hpccm.config.set_linux_distro('centos')
         elif self.__distro == 'centos7':
diff --git a/test/test_baseimage.py b/test/test_baseimage.py
index 4774f427..87362b52 100644
--- a/test/test_baseimage.py
+++ b/test/test_baseimage.py
@@ -188,6 +188,20 @@ def test_detect_ubuntu_20(self):
         self.assertEqual(hpccm.config.g_linux_distro, linux_distro.UBUNTU)
         self.assertEqual(hpccm.config.g_linux_version, Version('20.04'))
 
+    @docker
+    def test_detect_ubuntu_22(self):
+        """Base image Linux distribution detection"""
+        b = baseimage(image='ubuntu:22.04')
+        self.assertEqual(hpccm.config.g_linux_distro, linux_distro.UBUNTU)
+        self.assertEqual(hpccm.config.g_linux_version, Version('22.04'))
+
+    @docker
+    def test_detect_ubuntu_24(self):
+        """Base image Linux distribution detection"""
+        b = baseimage(image='ubuntu:24.04')
+        self.assertEqual(hpccm.config.g_linux_distro, linux_distro.UBUNTU)
+        self.assertEqual(hpccm.config.g_linux_version, Version('24.04'))
+
     @docker
     def test_detect_centos(self):
         """Base image Linux distribution detection"""
@@ -216,6 +230,13 @@ def test_detect_rockylinux_8(self):
         self.assertEqual(hpccm.config.g_linux_distro, linux_distro.ROCKYLINUX)
         self.assertEqual(hpccm.config.g_linux_version, Version('8.0'))
 
+    @docker
+    def test_detect_rockylinux_9(self):
+        """Base image Linux distribution detection"""
+        b = baseimage(image='rockylinux/rockylinux:9')
+        self.assertEqual(hpccm.config.g_linux_distro, linux_distro.ROCKYLINUX)
+        self.assertEqual(hpccm.config.g_linux_version, Version('9.0'))
+
     @docker
     def test_detect_ubi7(self):
         """Base image Linux distribution detection"""
@@ -274,6 +295,20 @@ def test_distro_ubuntu20(self):
         self.assertEqual(hpccm.config.g_linux_distro, linux_distro.UBUNTU)
         self.assertEqual(hpccm.config.g_linux_version, Version('20.04'))
 
+    @docker
+    def test_distro_ubuntu22(self):
+        """Base image Linux distribution specification"""
+        b = baseimage(image='foo', _distro='ubuntu22')
+        self.assertEqual(hpccm.config.g_linux_distro, linux_distro.UBUNTU)
+        self.assertEqual(hpccm.config.g_linux_version, Version('22.04'))
+
+    @docker
+    def test_distro_ubuntu24(self):
+        """Base image Linux distribution specification"""
+        b = baseimage(image='foo', _distro='ubuntu24')
+        self.assertEqual(hpccm.config.g_linux_distro, linux_distro.UBUNTU)
+        self.assertEqual(hpccm.config.g_linux_version, Version('24.04'))
+
     @docker
     def test_distro_centos(self):
         """Base image Linux distribution specification"""
@@ -295,6 +330,20 @@ def test_distro_centos8(self):
         self.assertEqual(hpccm.config.g_linux_distro, linux_distro.CENTOS)
         self.assertEqual(hpccm.config.g_linux_version, Version('8.0'))
 
+    @docker
+    def test_distro_rockylinux8(self):
+        """Base image Linux distribution specification"""
+        b = baseimage(image='foo', _distro='rockylinux8')
+        self.assertEqual(hpccm.config.g_linux_distro, linux_distro.CENTOS)
+        self.assertEqual(hpccm.config.g_linux_version, Version('8.0'))
+
+    @docker
+    def test_distro_rockylinux9(self):
+        """Base image Linux distribution specification"""
+        b = baseimage(image='foo', _distro='rockylinux9')
+        self.assertEqual(hpccm.config.g_linux_distro, linux_distro.CENTOS)
+        self.assertEqual(hpccm.config.g_linux_version, Version('9.0'))
+
     @docker
     def test_distro_nonexistent(self):
         """Base image Linux distribution specification"""