From 52c78481eb14659da8ccaa2f4c74bc1b8240572a Mon Sep 17 00:00:00 2001
From: Markus Rudy <mr@edgeless.systems>
Date: Mon, 16 Dec 2024 08:51:10 +0100
Subject: [PATCH] node-installer: add test for kata config

---
 nodeinstaller/internal/config/kata_runtime.go |  7 ++++
 .../internal/config/kata_runtime_test.go      | 40 +++++++++++++++++++
 nodeinstaller/node-installer.go               |  2 +-
 3 files changed, 48 insertions(+), 1 deletion(-)
 create mode 100644 nodeinstaller/internal/config/kata_runtime_test.go

diff --git a/nodeinstaller/internal/config/kata_runtime.go b/nodeinstaller/internal/config/kata_runtime.go
index f58e74aea4..25ffdead69 100644
--- a/nodeinstaller/internal/config/kata_runtime.go
+++ b/nodeinstaller/internal/config/kata_runtime.go
@@ -3,6 +3,8 @@
 
 package config
 
+import "github.com/pelletier/go-toml/v2"
+
 // KataRuntimeConfig is the configuration for the Kata runtime.
 // Source: https://github.com/kata-containers/kata-containers/blob/4029d154ba0c26fcf4a8f9371275f802e3ef522c/src/runtime/pkg/katautils/config.go
 // This is a simplified version of the actual configuration.
@@ -14,6 +16,11 @@ type KataRuntimeConfig struct {
 	Runtime    KataRuntime
 }
 
+// Marshal encodes the configuration as TOML.
+func (k *KataRuntimeConfig) Marshal() ([]byte, error) {
+	return toml.Marshal(k)
+}
+
 // Image is the configuration for the image.
 type Image map[string]any
 
diff --git a/nodeinstaller/internal/config/kata_runtime_test.go b/nodeinstaller/internal/config/kata_runtime_test.go
new file mode 100644
index 0000000000..25e81fee9e
--- /dev/null
+++ b/nodeinstaller/internal/config/kata_runtime_test.go
@@ -0,0 +1,40 @@
+// Copyright 2024 Edgeless Systems GmbH
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package config_test
+
+import (
+	"testing"
+
+	"github.com/edgelesssys/contrast/internal/platforms"
+	"github.com/edgelesssys/contrast/nodeinstaller/internal/constants"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestKataConfig(t *testing.T) {
+	// This is a regression test that ensures the `agent.kata` section is not optimized away. Empty
+	// section and no section are handled differently by Kata, so we make sure that this section is
+	// always present.
+	for _, platform := range platforms.All() {
+		t.Run(platform.String(), func(t *testing.T) {
+			require := require.New(t)
+			assert := assert.New(t)
+			cfg, err := constants.KataRuntimeConfig("/", platform, "", false)
+			require.NoError(err)
+			configBytes, err := cfg.Marshal()
+			require.NoError(err)
+			assert.Contains(string(configBytes), "[Agent.kata]")
+			assert.Contains(string(configBytes), "[Runtime]")
+
+			switch platform {
+			case platforms.K3sQEMUSNP, platforms.K3sQEMUTDX, platforms.MetalQEMUSNP, platforms.MetalQEMUTDX, platforms.RKE2QEMUTDX:
+				assert.Contains(string(configBytes), "[Hypervisor.qemu]")
+			case platforms.AKSCloudHypervisorSNP:
+				assert.Contains(string(configBytes), "[Hypervisor.clh]")
+			default:
+				assert.Fail("missing hypervisor test expectations")
+			}
+		})
+	}
+}
diff --git a/nodeinstaller/node-installer.go b/nodeinstaller/node-installer.go
index 158bbbca98..e54f0cb98b 100644
--- a/nodeinstaller/node-installer.go
+++ b/nodeinstaller/node-installer.go
@@ -182,7 +182,7 @@ func containerdRuntimeConfig(basePath, configPath string, platform platforms.Pla
 	if err != nil {
 		return fmt.Errorf("generating kata runtime config: %w", err)
 	}
-	rawConfig, err := toml.Marshal(kataRuntimeConfig)
+	rawConfig, err := kataRuntimeConfig.Marshal()
 	if err != nil {
 		return fmt.Errorf("marshaling kata runtime config: %w", err)
 	}