package service

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	"github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/runtime"
	"k8s.io/client-go/kubernetes/scheme"
	"sigs.k8s.io/controller-runtime/pkg/client"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"
)

func Test_podRestarterV1_Restart(t *testing.T) {
	tests := []struct {
		name   string
		pod    *corev1.Pod
		force  bool
		expect func(*WithT, client.Client, string)
	}{
		{
			name:  "graceful restart tidb pod",
			pod:   buildTestPod("tidb", nil),
			force: false,
			expect: func(g *WithT, c client.Client, name string) {
				var p corev1.Pod
				g.Expect(c.Get(context.TODO(), client.ObjectKey{
					Name: name,
				}, &p)).To(Succeed())
				g.Expect(p.Annotations).To(HaveKeyWithValue(v1alpha1.TiDBGracefulShutdownAnnKey, gracefulRestartAnnoVal))
			},
		},
		{
			name:  "force restart tikv pod",
			pod:   buildTestPod("tikv", nil),
			force: true,
			expect: func(g *WithT, c client.Client, name string) {
				// The pod should be deleted
				var p corev1.Pod
				g.Expect(c.Get(context.TODO(), client.ObjectKey{Name: name}, &p)).NotTo(Succeed())
			},
		},
		{
			name:  "graceful restart pd pod",
			pod:   buildTestPod("pd", nil),
			force: false,
			expect: func(g *WithT, c client.Client, name string) {
				var p corev1.Pod
				g.Expect(c.Get(context.TODO(), client.ObjectKey{Name: name}, &p)).To(Succeed())
				g.Expect(p.Annotations).To(HaveKeyWithValue(v1alpha1.PDLeaderTransferAnnKey, gracefulRestartAnnoVal))
			},
		},
		{
			name:  "restart tiflash pod",
			pod:   buildTestPod("tiflash", nil),
			force: false,
			expect: func(g *WithT, c client.Client, name string) {
				// The pod should be deleted
				var p corev1.Pod
				g.Expect(c.Get(context.TODO(), client.ObjectKey{Name: name}, &p)).NotTo(Succeed())
			},
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			g := NewGomegaWithT(t)
			s := runtime.NewScheme()
			g.Expect(scheme.AddToScheme(s)).To(Succeed())
			cli := fake.NewClientBuilder().WithScheme(s).WithObjects(tt.pod).Build()

			p := &podRestarterV1{}
			g.Expect(p.Restart(context.TODO(), cli, tt.pod, tt.force)).To(Succeed())
			tt.expect(g, cli, tt.pod.Name)
		})
	}
}