From 23bba03bbd1d4f94b1486939f9ebd06b73576382 Mon Sep 17 00:00:00 2001
From: Michael Vogt <michael.vogt@gmail.com>
Date: Fri, 16 Aug 2024 10:25:31 +0200
Subject: [PATCH] manifest: extract tomlPkgsFor() helper and fix ordering

The current code will install python3-toml on fedora and also
on rhel11 (once it comes out) because that is the default for
unhandled distros.

This commit makes the various toml libs for each distro more
explicit, see also https://github.com/osbuild/osbuild/pull/1851
---
 pkg/manifest/os.go      | 38 ++++++++++++++++++++----------------
 pkg/manifest/os_test.go | 43 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 64 insertions(+), 17 deletions(-)

diff --git a/pkg/manifest/os.go b/pkg/manifest/os.go
index 88fcd039b7..38554595f2 100644
--- a/pkg/manifest/os.go
+++ b/pkg/manifest/os.go
@@ -163,7 +163,6 @@ type OS struct {
 	// OSTreeParent source spec (optional). If nil the new commit (if
 	// applicable) will have no parent
 	OSTreeParent *ostree.SourceSpec
-
 	// Enabling Bootupd runs bootupctl generate-update-metadata in the tree to
 	// transform /usr/lib/ostree-boot into a bootupd-compatible update
 	// payload. Only works with ostree-based images.
@@ -289,6 +288,25 @@ func (p *OS) getContainerSources() []container.SourceSpec {
 	return p.OSCustomizations.Containers
 }
 
+func tomlPkgsFor(distro Distro) []string {
+	switch distro {
+	case DISTRO_EL7:
+		// nothing needs toml in rhel7
+		panic("no support for toml on rhel7")
+	case DISTRO_EL8:
+		// deprecated, needed for backwards compatibility (EL8 manifests)
+		return []string{"python3-pytoml"}
+	case DISTRO_EL9:
+		// older unmaintained lib, needed for backwards compatibility
+		return []string{"python3-toml"}
+	default:
+		// No extra package needed for reading, on rhel10 and
+		// fedora as stdlib has "tomlib" but we need tomli-w
+		// for writing
+		return []string{"python3-tomli-w"}
+	}
+}
+
 func (p *OS) getBuildPackages(distro Distro) []string {
 	packages := p.platform.GetBuildPackages()
 	if p.PartitionTable != nil {
@@ -315,14 +333,7 @@ func (p *OS) getBuildPackages(distro Distro) []string {
 
 	if len(p.OSCustomizations.Containers) > 0 {
 		if p.OSCustomizations.ContainersStorage != nil {
-			switch distro {
-			case DISTRO_EL8:
-				packages = append(packages, "python3-pytoml")
-			case DISTRO_EL10:
-				// no extra package needed, stdlib has "tomli"
-			default:
-				packages = append(packages, "python3-toml")
-			}
+			packages = append(packages, tomlPkgsFor(distro)...)
 		}
 		packages = append(packages, "skopeo")
 	}
@@ -332,14 +343,7 @@ func (p *OS) getBuildPackages(distro Distro) []string {
 	}
 
 	if p.BootcConfig != nil {
-		switch distro {
-		case DISTRO_EL8:
-			packages = append(packages, "python3-pytoml")
-		case DISTRO_EL10:
-			// no extra package needed, stdlib has "tomli"
-		default:
-			packages = append(packages, "python3-toml")
-		}
+		packages = append(packages, tomlPkgsFor(distro)...)
 	}
 
 	return packages
diff --git a/pkg/manifest/os_test.go b/pkg/manifest/os_test.go
index be229ebc8e..dddb9aeada 100644
--- a/pkg/manifest/os_test.go
+++ b/pkg/manifest/os_test.go
@@ -4,6 +4,9 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/osbuild/images/internal/common"
+	"github.com/osbuild/images/pkg/container"
+	"github.com/osbuild/images/pkg/customizations/bootc"
 	"github.com/osbuild/images/pkg/customizations/subscription"
 	"github.com/osbuild/images/pkg/osbuild"
 	"github.com/osbuild/images/pkg/platform"
@@ -181,3 +184,43 @@ func TestBootupdStage(t *testing.T) {
 	st := findStage("org.osbuild.bootupd.gen-metadata", pipeline.Stages)
 	require.NotNil(t, st)
 }
+
+func TestTomlLibUsedNoneByDefault(t *testing.T) {
+	os := NewTestOS()
+	buildPkgs := os.getBuildPackages(DISTRO_FEDORA)
+	for _, pkg := range []string{"python3-pytoml", "python3-toml", "python3-tomli-w"} {
+		assert.NotContains(t, buildPkgs, pkg)
+	}
+}
+
+func TestTomlLibUsedForContainer(t *testing.T) {
+	os := NewTestOS()
+	os.OSCustomizations.Containers = []container.SourceSpec{
+		{Source: "some-source"},
+	}
+	os.OSCustomizations.ContainersStorage = common.ToPtr("foo")
+
+	testTomlPkgsFor(t, os)
+}
+
+func TestTomlLibUsedForBootcConfig(t *testing.T) {
+	os := NewTestOS()
+	os.BootcConfig = &bootc.Config{Filename: "something"}
+
+	testTomlPkgsFor(t, os)
+}
+
+func testTomlPkgsFor(t *testing.T, os *OS) {
+	for _, tc := range []struct {
+		distro          Distro
+		expectedTomlPkg string
+	}{
+		{DISTRO_EL8, "python3-pytoml"},
+		{DISTRO_EL9, "python3-toml"},
+		{DISTRO_EL10, "python3-tomli-w"},
+		{DISTRO_FEDORA, "python3-tomli-w"},
+	} {
+		buildPkgs := os.getBuildPackages(tc.distro)
+		assert.Contains(t, buildPkgs, tc.expectedTomlPkg)
+	}
+}