package gardener

import (
	"context"
	"testing"

	"github.com/crossplane/crossplane-runtime/pkg/test"
	"github.com/google/go-cmp/cmp"
	. "github.com/onsi/gomega"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	cloudv1 "github.com/tidbcloud/infra-api/api/v1"
	gardencorev1beta1 "github.com/tidbcloud/infra-api/pkg/gardener/core/apis/v1beta1"
	"github.com/tidbcloud/infra-api/pkg/nodegroup"
)

func TestApplyNodeGroupSet(t *testing.T) {
	cases := map[string]struct {
		ngs     *nodegroup.Set
		current []gardencorev1beta1.Worker
		desired []gardencorev1beta1.Worker
		policy  NodeGroupPolicy
	}{
		"ApplyNewNodeGroupSet": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb1"},
					{Name: "tidb2"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"CleanupNodeGroupSet": {
			ngs: &nodegroup.Set{
				Name:       "cluster1",
				NodeGroups: nil,
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"ScaleOutNodeGroupSet": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb1"},
					{Name: "tidb2"},
					{Name: "tidb3"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb3",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"ScaleInNodeGroupSet": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb1"},
					{Name: "tidb2"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb3",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"RollingUpdateNodeGroupSet": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{
						Name:    "tidb1",
						Minimum: 2,
					},
					{
						Name:    "tidb2",
						Minimum: 2,
					},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
				},
				{
					Name:    "tidb2",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
				},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 2,
				},
				{
					Name:    "tidb2",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 2,
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"ReplaceNodeGroupSet": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tikv1"},
					{Name: "tikv2"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tikv1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tikv2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"AdoptOrphanNodeGroups": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb1"},
					{Name: "tidb2"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{Name: "tidb1"},
				{Name: "tidb2"},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"AdoptOrphanNodeGroupsWithScaleIn": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb1"},
					{Name: "tidb2"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{Name: "tidb1"},
				{Name: "tidb2"},
				{Name: "tidb3"},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				// tidb3 is leaked if there is a pending scale-in operation, which is not likely to happen
				{Name: "tidb3"},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"UpdateNodeGroupShouldKeepDefaults": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{
						Name:    "tidb1",
						Minimum: 2,
						MachineSpec: &cloudv1.MachineSpec{
							MachineImage: "centos",
						},
					},
					{
						Name:    "tidb2",
						Minimum: 2,
						MachineSpec: &cloudv1.MachineSpec{
							MachineImage: "centos",
						},
					},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name:    "centos",
							Version: pointer.String("test"),
						},
					},
					Minimum: 1,
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name:    "centos",
							Version: pointer.String("test"),
						},
					},
					Minimum: 1,
				},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{Name: "master2"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name:    "centos",
							Version: pointer.String("test"),
						},
					},
					Minimum: 2,
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name:    "centos",
							Version: pointer.String("test"),
						},
					},
					Minimum: 2,
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"RollingUpdateMachineTypeWhenScale": {
			current: []gardencorev1beta1.Worker{
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
					Minimum: 1,
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{
						Name: "tidb1",
						MachineSpec: &cloudv1.MachineSpec{
							MachineImage: "large",
						},
						Minimum: 2,
					},
				},
			},
			desired: []gardencorev1beta1.Worker{
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "large",
						},
					},
					Minimum: 2,
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"ScaleOnlyShouldNotMutateMachineType": {
			current: []gardencorev1beta1.Worker{
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
					Minimum: 1,
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{
						Name: "tidb1",
						MachineSpec: &cloudv1.MachineSpec{
							MachineImage: "large",
						},
						Minimum: 2,
					},
				},
			},
			desired: []gardencorev1beta1.Worker{
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
					Minimum: 2,
				},
			},
			policy: NodeGroupPolicyScaleOnly,
		},
		"OrderShouldBeStable": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb3"},
					{Name: "tidb2"},
					{Name: "tidb1"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb2",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:   "tidb3",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"Dryrun": {
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb3"},
					{Name: "tidb2"},
					{Name: "tidb1"},
				},
			},
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
			},
			desired: []gardencorev1beta1.Worker{
				{Name: "master1"},
			},
			policy: NodeGroupPolicyDryrun,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			ngs := tc.current
			g := NewGomegaWithT(t)
			gardenEnv := &GardEnvironment{
				client: &test.MockClient{
					MockGet: func(ctx context.Context, key client.ObjectKey, obj client.Object) error {
						if shoot, ok := obj.(*gardencorev1beta1.Shoot); ok {
							shoot.Spec = gardencorev1beta1.ShootSpec{
								Provider: gardencorev1beta1.Provider{
									Workers: ngs,
								},
							}
						}
						return nil
					},
					MockUpdate: func(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
						if shoot, ok := obj.(*gardencorev1beta1.Shoot); ok {
							ngs = shoot.Spec.Provider.Workers
						}
						return nil
					},
				},
			}
			ctx := context.TODO()
			_, err := gardenEnv.ApplyNodeGroupSet(ctx, "dummy", tc.ngs, tc.policy)
			g.Expect(err).To(Succeed())
			if diff := cmp.Diff(tc.desired, ngs); diff != "" {
				t.Errorf("diff:\n%s", diff)
			}
		})
	}
}

func TestChangesOfApplyNodeGroupSet(t *testing.T) {
	cases := map[string]struct {
		current []gardencorev1beta1.Worker
		ngs     *nodegroup.Set
		policy  NodeGroupPolicy
		changes []NodeGroupChange
	}{
		"AddNodeGroup": {
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "tidb1"},
				},
			},
			changes: []NodeGroupChange{{
				Operation: NodeGroupOpCreate,
				After: &gardencorev1beta1.Worker{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			}},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"DeleteNodeGroup": {
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
			},
			changes: []NodeGroupChange{{
				Operation: NodeGroupOpDelete,
				Before: &gardencorev1beta1.Worker{
					Name:   "tidb1",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			}},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"UpdateNodeGroup": {
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{
						Name:    "tidb1",
						Minimum: 2,
						MachineSpec: &cloudv1.MachineSpec{
							MachineImage: "large",
						},
					},
				},
			},
			changes: []NodeGroupChange{{
				Operation: NodeGroupOpUpdate,
				Before: &gardencorev1beta1.Worker{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
				},
				After: &gardencorev1beta1.Worker{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 2,
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "large",
						},
					},
				},
			}},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"NoOperation": {
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{
						Name:    "tidb1",
						Minimum: 1,
						MachineSpec: &cloudv1.MachineSpec{
							MachineImage: "small",
						},
					},
				},
			},
			changes: nil,
			policy:  NodeGroupPolicyRollingUpdate,
		},
		"ScaleNodeGroup": {
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{
						Name:    "tidb1",
						Minimum: 2,
						MachineSpec: &cloudv1.MachineSpec{
							MachineImage: "large",
						},
					},
				},
			},
			changes: []NodeGroupChange{{
				Operation: NodeGroupOpUpdate,
				Before: &gardencorev1beta1.Worker{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
				},
				After: &gardencorev1beta1.Worker{
					Name:    "tidb1",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 2,
					Machine: gardencorev1beta1.Machine{
						Image: &gardencorev1beta1.ShootMachineImage{
							Name: "small",
						},
					},
				},
			}},
			policy: NodeGroupPolicyScaleOnly,
		},
		"Comprehensive": {
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{
					Name:   "ToDelete",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:    "ToUpdate",
					Minimum: 1,
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:    "NoChange",
					Minimum: 1,
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "ToCreate"},
					{
						Name:    "ToUpdate",
						Minimum: 2,
					},
					{
						Name:    "NoChange",
						Minimum: 1,
					},
				},
			},
			changes: []NodeGroupChange{{
				Operation: NodeGroupOpDelete,
				Before: &gardencorev1beta1.Worker{
					Name:   "ToDelete",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			}, {
				Operation: NodeGroupOpCreate,
				After: &gardencorev1beta1.Worker{
					Name:   "ToCreate",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			}, {
				Operation: NodeGroupOpUpdate,
				Before: &gardencorev1beta1.Worker{
					Name:    "ToUpdate",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
				},
				After: &gardencorev1beta1.Worker{
					Name:    "ToUpdate",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 2,
				},
			}},
			policy: NodeGroupPolicyRollingUpdate,
		},
		"Dryrun": {
			current: []gardencorev1beta1.Worker{
				{Name: "master1"},
				{
					Name:   "ToDelete",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:    "ToUpdate",
					Minimum: 1,
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
				{
					Name:    "NoChange",
					Minimum: 1,
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			},
			ngs: &nodegroup.Set{
				Name: "cluster1",
				NodeGroups: []cloudv1.NodegroupSpec{
					{Name: "ToCreate"},
					{
						Name:    "ToUpdate",
						Minimum: 2,
					},
					{
						Name:    "NoChange",
						Minimum: 1,
					},
				},
			},
			changes: []NodeGroupChange{{
				Operation: NodeGroupOpDelete,
				Before: &gardencorev1beta1.Worker{
					Name:   "ToDelete",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			}, {
				Operation: NodeGroupOpCreate,
				After: &gardencorev1beta1.Worker{
					Name:   "ToCreate",
					Labels: map[string]string{NodeGroupSetOwnerKey: "cluster1"},
				},
			}, {
				Operation: NodeGroupOpUpdate,
				Before: &gardencorev1beta1.Worker{
					Name:    "ToUpdate",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 1,
				},
				After: &gardencorev1beta1.Worker{
					Name:    "ToUpdate",
					Labels:  map[string]string{NodeGroupSetOwnerKey: "cluster1"},
					Minimum: 2,
				},
			}},
			policy: NodeGroupPolicyDryrun,
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			gardenEnv := &GardEnvironment{
				client: &test.MockClient{
					MockGet: func(ctx context.Context, key client.ObjectKey, obj client.Object) error {
						if shoot, ok := obj.(*gardencorev1beta1.Shoot); ok {
							shoot.Spec = gardencorev1beta1.ShootSpec{
								Provider: gardencorev1beta1.Provider{
									Workers: tc.current,
								},
							}
						}
						return nil
					},
					MockUpdate: func(ctx context.Context, obj client.Object, opts ...client.UpdateOption) error {
						return nil
					},
				},
			}
			ctx := context.TODO()
			changes, err := gardenEnv.ApplyNodeGroupSet(ctx, "dummy", tc.ngs, tc.policy)
			g.Expect(err).To(Succeed())
			g.Expect(changes).To(ConsistOf(tc.changes))
		})
	}
}

func TestUpdateShootFirewall(t *testing.T) {
	cases := map[string]struct {
		shoot  *gardencorev1beta1.Shoot
		cidrs  []string
		expect *gardencorev1beta1.Shoot
	}{
		"SetCidrs": {
			shoot: &gardencorev1beta1.Shoot{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "garden-test",
				},
			},
			cidrs: []string{
				"8.8.8.8/32",
			},
			expect: &gardencorev1beta1.Shoot{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "garden-test",
				},
				Spec: gardencorev1beta1.ShootSpec{
					LoadBalancerSourceRanges: []string{
						"8.8.8.8/32",
					},
				},
			},
		},
		"UpdateCidrs": {
			shoot: &gardencorev1beta1.Shoot{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "garden-test",
				},
				Spec: gardencorev1beta1.ShootSpec{
					LoadBalancerSourceRanges: []string{
						"0.0.0.0/0",
						"7.7.7.7/32",
					},
				},
			},
			cidrs: []string{
				"6.6.6.6/32",
				"8.8.8.8/32",
			},
			expect: &gardencorev1beta1.Shoot{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "garden-test",
				},
				Spec: gardencorev1beta1.ShootSpec{
					LoadBalancerSourceRanges: []string{
						"6.6.6.6/32",
						"8.8.8.8/32",
					},
				},
			},
		},
		"UpdateCidrsShouldNotMutateOtherFields": {
			shoot: &gardencorev1beta1.Shoot{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "garden-test",
				},
				Spec: gardencorev1beta1.ShootSpec{
					LoadBalancerSourceRanges: []string{
						"0.0.0.0/0",
					},
					Provider: gardencorev1beta1.Provider{
						Workers: []gardencorev1beta1.Worker{{
							Name: "worker-1",
						}},
					},
				},
			},
			cidrs: []string{
				"8.8.8.8/32",
			},
			expect: &gardencorev1beta1.Shoot{
				ObjectMeta: metav1.ObjectMeta{
					Name:      "test",
					Namespace: "garden-test",
				},
				Spec: gardencorev1beta1.ShootSpec{
					LoadBalancerSourceRanges: []string{
						"8.8.8.8/32",
					},
					Provider: gardencorev1beta1.Provider{
						Workers: []gardencorev1beta1.Worker{{
							Name: "worker-1",
						}},
					},
				},
			},
		},
	}
	for name, tc := range cases {
		t.Run(name, func(t *testing.T) {
			ctx := context.Background()
			g := NewGomegaWithT(t)
			scheme := runtime.NewScheme()
			g.Expect(gardencorev1beta1.AddToScheme(scheme)).To(Succeed())
			cli := fake.NewClientBuilder().WithScheme(scheme).WithObjects(tc.shoot).Build()
			gardEnv := &GardEnvironment{client: cli, projectName: "test"}
			err := gardEnv.UpdateShootFirewall(ctx, tc.shoot.Name, tc.cidrs)
			g.Expect(err).To(Succeed())
			updated := &gardencorev1beta1.Shoot{}
			g.Expect(cli.Get(ctx, types.NamespacedName{Namespace: tc.shoot.Namespace, Name: tc.shoot.Name}, updated)).To(Succeed())
			g.Expect(updated.Spec).To(Equal(tc.expect.Spec))
		})
	}
}