package nodedrainer

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/go-logr/logr"
	. "github.com/onsi/gomega"
	"github.com/tidbcloud/aws-shared-provider/pkg/scheme"
	"github.com/tidbcloud/aws-shared-provider/pkg/types"
	"github.com/tidbcloud/aws-shared-provider/pkg/utils/labels"
	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/tools/record"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
	"sigs.k8s.io/controller-runtime/pkg/log"
)

func TestNodeDrainer_do(t *testing.T) {
	now := metav1.Now()
	tests := []struct {
		name     string
		nodeList []runtime.Object
		podList  []runtime.Object

		wantDrainedNodes   []string
		wantUnDrainedNodes []string
		wantPods           []string
	}{
		{
			name: "no valid nodes",
			nodeList: []runtime.Object{
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_not_label",
					},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_not_on_demand",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey: "tidb-xxx",
						},
					},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_terminating",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						DeletionTimestamp: &now,
						Finalizers:        []string{"fake-node-finalizer"},
					},
				},
			},
			wantUnDrainedNodes: []string{"node_not_on_demand", "node_terminating"},
		},
		{
			name: "test drain in different node group",
			nodeList: []runtime.Object{
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node0_drain",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx0",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node1_drain",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx1",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
				},
			},
			podList: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "not_evict_pod_for_terminating",
						Namespace: "tidb-nodedrainer-1",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
						DeletionTimestamp: &now,
						Finalizers:        []string{types.TidbPoolFinalizer},
					},
					Spec: corev1.PodSpec{NodeName: "node1_drain"},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "evict_pod_node0",
						Namespace: "tidb-pool-0",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
					},
					Spec: corev1.PodSpec{NodeName: "node0_drain"},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name:      "evict_pod_node1",
						Namespace: "tidb-nodedrainer-1",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
					},
					Spec: corev1.PodSpec{NodeName: "node1_drain"},
				},
			},
			wantDrainedNodes: []string{"node0_drain", "node1_drain"},
			wantPods:         []string{"not_evict_pod_for_terminating"},
		},
		{
			name: "test drain node, but no evict any pods, because of found NotSafeToEvictPod pod",
			nodeList: []runtime.Object{
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_drain_no_evict_any_pods",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
				},
			},
			podList: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "can_evict_pod",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
					},
					Spec: corev1.PodSpec{NodeName: "node_drain_no_evict_any_pods"},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "skip_evict_pod",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
						Annotations: map[string]string{types.CASafeToEvict: "false"},
					},
					Spec: corev1.PodSpec{NodeName: "node_drain_no_evict_any_pods"},
				},
			},
			wantDrainedNodes: []string{"node_drain_no_evict_any_pods"},
			wantPods:         []string{"can_evict_pod", "skip_evict_pod"},
		},
		{
			name: "test drain all nodes",
			nodeList: []runtime.Object{
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node0_unschedule_drain_empty", // node0 can be first listed to drain
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
					Spec: corev1.NodeSpec{Unschedulable: true},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node1_drained_before", // node1 can be next listed to drain
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						Annotations: map[string]string{
							types.AnnoNodeDrainTimestamp: now.Add(time.Minute).Format(time.RFC3339),
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
					Spec: corev1.NodeSpec{Unschedulable: true},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node2_only_unschedule_drain",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
				},
			},
			podList: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "evict_pod_in_node1_drained_before",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
					},
					Spec: corev1.PodSpec{NodeName: "node1_drained_before"},
				},
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "not_evict_pod_evicted_before",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
					},
					Spec: corev1.PodSpec{NodeName: "node2_only_unschedule_drain"},
				},
			},
			wantDrainedNodes: []string{"node0_unschedule_drain_empty", "node1_drained_before", "node2_only_unschedule_drain"},
			wantPods:         []string{"not_evict_pod_evicted_before"},
		},
		{
			name: "test recover node_recover/node1_recover",
			nodeList: []runtime.Object{
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_recover",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						Annotations: map[string]string{
							types.AnnoNodeDrainTimestamp: now.Add(time.Minute).Format(time.RFC3339),
						},
						CreationTimestamp: now,
					},
					Spec: corev1.NodeSpec{Unschedulable: true},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_no_drain_timestamp", // keep Unschedulable field
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						CreationTimestamp: now,
					},
					Spec: corev1.NodeSpec{Unschedulable: true},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node1_recover",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						Annotations: map[string]string{
							types.AnnoNodeDrainTimestamp: now.Add(time.Minute).Format(time.RFC3339),
						},
						CreationTimestamp: now,
					},
					Spec: corev1.NodeSpec{Unschedulable: true},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_can_drain",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
				},
			},
			podList: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "in_node_can_drain",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
					},
					Spec: corev1.PodSpec{NodeName: "node_can_drain"},
				},
			},
			wantUnDrainedNodes: []string{"node_recover", "node_no_drain_timestamp", "node1_recover", "node_can_drain"},
			wantPods:           []string{"in_node_can_drain"},
		},
		{
			name: "test recover node_recover because of prevent-ca-scale-up",
			nodeList: []runtime.Object{
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_spot_prevent-ca-scale-up",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeSpot,
							labels.ServerlessComponentNodeLabelKey: "tidb",
							corev1.LabelTopologyZone:               "us-east-1a",
						},
						CreationTimestamp: now,
					},
					Spec: corev1.NodeSpec{
						Taints: []corev1.Taint{{
							Key:    types.PreventCAScaleUpTaintKey,
							Value:  fmt.Sprint(now.Unix()),
							Effect: corev1.TaintEffectNoSchedule,
						}},
					},
				},
				&corev1.Node{
					ObjectMeta: metav1.ObjectMeta{
						Name: "node_recover",
						Labels: map[string]string{
							labels.ServerlessNodeGroupLabelKey:     "tidb-od-xxx",
							labels.NodeDrainerCapacityTypeLabelKey: labels.NodeDrainerCapacityTypeOnDemand,
							labels.ServerlessComponentNodeLabelKey: "tidb",
							corev1.LabelTopologyZone:               "us-east-1a",
						},
						Annotations: map[string]string{
							types.AnnoNodeDrainTimestamp: now.Add(time.Minute).Format(time.RFC3339),
						},
						CreationTimestamp: metav1.NewTime(now.Add(-activeOnDemandNodeDuration * 2)),
					},
					Spec: corev1.NodeSpec{Unschedulable: true},
				},
			},
			podList: []runtime.Object{
				&corev1.Pod{
					ObjectMeta: metav1.ObjectMeta{
						Name: "in_node_recover",
						Labels: map[string]string{
							labels.NameLabelKey:      labels.TidbPodNameLabelValue,
							labels.ComponentLabelKey: labels.TidbPodComponentLabelValue,
						},
					},
					Spec: corev1.PodSpec{NodeName: "node_recover"},
				},
			},
			wantUnDrainedNodes: []string{"node_recover"},
			wantPods:           []string{"in_node_recover"},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)

			nd := &NodeDrainer{
				Logger: logr.New(log.NullLogSink{}),
				ctx:    context.Background(),
				client: fake.NewClientBuilder().WithScheme(scheme.Scheme).
					WithRuntimeObjects(tt.nodeList...).WithRuntimeObjects(tt.podList...).
					WithIndex(&corev1.Pod{}, "spec.nodeName", func(obj client.Object) []string {
						pod := obj.(*corev1.Pod)
						return []string{pod.Spec.NodeName}
					}).
					Build(),
				recorder: record.NewFakeRecorder(4096),
			}

			nd.do()

			for _, wantDrainedNode := range tt.wantDrainedNodes {
				node := &corev1.Node{}
				g.Expect(nd.client.Get(context.Background(), client.ObjectKey{Name: wantDrainedNode}, node)).To(Succeed())
				g.Expect(node.Annotations[types.AnnoNodeDrainTimestamp]).NotTo(BeEmpty())
				g.Expect(node.Spec.Unschedulable).To(BeTrue())
			}
			for _, wantUnDrainedNode := range tt.wantUnDrainedNodes {
				node := &corev1.Node{}
				g.Expect(nd.client.Get(context.Background(), client.ObjectKey{Name: wantUnDrainedNode}, node)).To(Succeed())
				g.Expect(node.Annotations).NotTo(HaveKey(types.AnnoNodeDrainTimestamp))
				for _, orgNode := range tt.nodeList {
					if orgNode.(*corev1.Node).Name == wantUnDrainedNode {
						if orgNode.(*corev1.Node).Annotations[types.AnnoNodeDrainTimestamp] == "" {
							g.Expect(node.Spec.Unschedulable).To(Equal(orgNode.(*corev1.Node).Spec.Unschedulable),
								"node %s should keep Unschedulable field", wantUnDrainedNode)
						} else {
							g.Expect(node.Spec.Unschedulable).To(BeFalse(),
								"node %s should be schedulable", wantUnDrainedNode)
						}
					}
				}
			}
			podList := &corev1.PodList{}
			g.Expect(nd.client.List(context.Background(), podList)).To(Succeed())
			g.Expect(len(podList.Items)).To(Equal(len(tt.wantPods)))
			for _, pod := range podList.Items {
				g.Expect(tt.wantPods).To(ContainElement(pod.Name))
			}
		})
	}
}