//go:build unit || all
// +build unit all

package eks

import (
	"encoding/json"
	"testing"

	. "github.com/onsi/gomega"
	appsv1 "k8s.io/api/apps/v1"
	corev1 "k8s.io/api/core/v1"
)

func TestMutateDaemonSetEnv(t *testing.T) {
	cases := map[string]struct {
		ds            *appsv1.DaemonSet
		containerName string
		envVar        corev1.EnvVar
		expected      *appsv1.DaemonSet
	}{
		"container not found": {
			ds: &appsv1.DaemonSet{
				Spec: appsv1.DaemonSetSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name: cniContainerName,
									Env: []corev1.EnvVar{
										{
											Name:  "ENABLE_POD_ENI",
											Value: "false",
										},
										{
											Name:  "ENABLE_WHATEVER",
											Value: "false",
										},
									},
								},
							},
						},
					},
				},
			},
			containerName: "non-exist",
			envVar:        corev1.EnvVar{Name: "ENABLE_WHATEVER", Value: "true"},
			// no change at all
			expected: &appsv1.DaemonSet{
				Spec: appsv1.DaemonSetSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name: cniContainerName,
									Env: []corev1.EnvVar{
										{
											Name:  "ENABLE_POD_ENI",
											Value: "false",
										},
										{
											Name:  "ENABLE_WHATEVER",
											Value: "false",
										},
									},
								},
							},
						},
					},
				},
			},
		},
		"container exist, env not found": {
			ds: &appsv1.DaemonSet{
				Spec: appsv1.DaemonSetSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name: cniContainerName,
									Env: []corev1.EnvVar{
										{
											Name:  "ENABLE_POD_ENI",
											Value: "false",
										},
										{
											Name:  "ENABLE_WHATEVER",
											Value: "false",
										},
									},
								},
							},
						},
					},
				},
			},
			containerName: cniContainerName,
			envVar:        corev1.EnvVar{Name: "NOT_EXIST", Value: "true"},
			// env should be appended
			expected: &appsv1.DaemonSet{
				Spec: appsv1.DaemonSetSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name: cniContainerName,
									Env: []corev1.EnvVar{
										{
											Name:  "ENABLE_POD_ENI",
											Value: "false",
										},
										{
											Name:  "ENABLE_WHATEVER",
											Value: "false",
										},
										{
											Name:  "NOT_EXIST",
											Value: "true",
										},
									},
								},
							},
						},
					},
				},
			},
		},
		"continer exist and env found": {
			ds: &appsv1.DaemonSet{
				Spec: appsv1.DaemonSetSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name: cniContainerName,
									Env: []corev1.EnvVar{
										{
											Name:  "ENABLE_POD_ENI",
											Value: "false",
										},
										{
											Name:  "ENABLE_WHATEVER",
											Value: "false",
										},
									},
								},
							},
						},
					},
				},
			},
			containerName: cniContainerName,
			envVar:        corev1.EnvVar{Name: "ENABLE_WHATEVER", Value: "true"},
			// env should be modified
			expected: &appsv1.DaemonSet{
				Spec: appsv1.DaemonSetSpec{
					Template: corev1.PodTemplateSpec{
						Spec: corev1.PodSpec{
							Containers: []corev1.Container{
								{
									Name: cniContainerName,
									Env: []corev1.EnvVar{
										{
											Name:  "ENABLE_POD_ENI",
											Value: "false",
										},
										{
											Name:  "ENABLE_WHATEVER",
											Value: "true",
										},
									},
								},
							},
						},
					},
				},
			},
		},
	}

	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			mutateDaemonSetEnv(tc.ds, tc.containerName, tc.envVar)
			str, err := json.Marshal(tc.ds)
			g.Expect(err).To(Succeed())
			expected, err := json.Marshal(tc.expected)
			g.Expect(err).To(Succeed())
			g.Expect(str).Should(Equal(expected))
		})
	}

}