package gardener

import (
	"context"
	"testing"

	. "github.com/onsi/gomega"
	corev1 "k8s.io/api/core/v1"
	storagev1 "k8s.io/api/storage/v1"
	"k8s.io/apimachinery/pkg/types"
	"k8s.io/utils/pointer"
	"sigs.k8s.io/controller-runtime/pkg/client/fake"

	"github.com/tidbcloud/infra-api/pkg/k8s"
)

func TestAWSEnsureStorageClasses(t *testing.T) {
	g := NewGomegaWithT(t)
	ctx := context.Background()

	client := fake.NewClientBuilder().Build()
	err := ensureStorageClasses(ctx, client, "aws", "v1.18")
	g.Expect(err).Should(Succeed())

	// check remote storage
	sc := new(storagev1.StorageClass)
	err = client.Get(ctx, types.NamespacedName{Name: k8s.DefaultStorageClassName}, sc)
	g.Expect(err).Should(Succeed())

	waitForFirstConsumer := storagev1.VolumeBindingWaitForFirstConsumer
	reclaimDelete := corev1.PersistentVolumeReclaimDelete
	expect := &storagev1.StorageClass{
		VolumeBindingMode:    &waitForFirstConsumer,
		AllowVolumeExpansion: pointer.BoolPtr(true),
		ReclaimPolicy:        &reclaimDelete,
		Provisioner:          "kubernetes.io/aws-ebs",
		Parameters: map[string]string{
			"type":   "gp2",
			"fsType": "ext4",
		},
		MountOptions: []string{"nodelalloc,noatime"},
	}
	sc.ObjectMeta = expect.ObjectMeta
	sc.TypeMeta = expect.TypeMeta
	g.Expect(expect).Should(Equal(sc))

	// check local storage
	sc = new(storagev1.StorageClass)
	err = client.Get(ctx, types.NamespacedName{Name: "local-storage"}, sc)
	g.Expect(err).Should(Succeed())
	expect = &storagev1.StorageClass{
		Provisioner:       "kubernetes.io/no-provisioner",
		VolumeBindingMode: &waitForFirstConsumer,
		MountOptions:      []string{"nodelalloc,noatime"},
	}
	sc.ObjectMeta = expect.ObjectMeta
	sc.TypeMeta = expect.TypeMeta
	g.Expect(expect).Should(Equal(sc))

	// check gp3
	sc = new(storagev1.StorageClass)
	err = client.Get(ctx, types.NamespacedName{Name: k8s.TiKVGp3StorageClassName}, sc)
	g.Expect(err).Should(Succeed())
	expect = &storagev1.StorageClass{
		VolumeBindingMode:    &waitForFirstConsumer,
		AllowVolumeExpansion: pointer.BoolPtr(true),
		ReclaimPolicy:        &reclaimDelete,
		Provisioner:          "ebs.csi.aws.com",
		Parameters: map[string]string{
			"csi.storage.k8s.io/fstype": "ext4",
			"type":                      "gp3",
			"iops":                      k8s.TiKVGp3IOPS,
			"throughput":                k8s.TiKVGp3ThroughputMBps,
			"encrypted":                 "true",
		},
		MountOptions: []string{"nodelalloc,noatime"},
	}
	sc.ObjectMeta = expect.ObjectMeta
	sc.TypeMeta = expect.TypeMeta
	g.Expect(expect).Should(Equal(sc))

	sc = new(storagev1.StorageClass)
	err = client.Get(ctx, types.NamespacedName{Name: k8s.TiFlashGp3StorageClassName}, sc)
	g.Expect(err).Should(Succeed())
	expect = &storagev1.StorageClass{
		VolumeBindingMode:    &waitForFirstConsumer,
		AllowVolumeExpansion: pointer.BoolPtr(true),
		ReclaimPolicy:        &reclaimDelete,
		Provisioner:          "ebs.csi.aws.com",
		Parameters: map[string]string{
			"csi.storage.k8s.io/fstype": "ext4",
			"type":                      "gp3",
			"iops":                      k8s.TiFlashGp3IOPS,
			"throughput":                k8s.TiFlashGp3ThroughputMBps,
			"encrypted":                 "true",
		},
		MountOptions: []string{"nodelalloc,noatime"},
	}
	sc.ObjectMeta = expect.ObjectMeta
	sc.TypeMeta = expect.TypeMeta
	g.Expect(expect).Should(Equal(sc))

	// just check duplicate call will not cause error
	err = ensureStorageClasses(ctx, client, "aws", "v1.18")
	g.Expect(err).Should(Succeed())
}

func TestGCPEnsureStorageClasses(t *testing.T) {
	g := NewGomegaWithT(t)
	ctx := context.Background()

	client := fake.NewClientBuilder().Build()
	err := ensureStorageClasses(ctx, client, "gcp", "v1.18")
	g.Expect(err).Should(Succeed())

	// check remote storage
	sc := new(storagev1.StorageClass)
	err = client.Get(ctx, types.NamespacedName{Name: k8s.DefaultStorageClassName}, sc)
	g.Expect(err).Should(Succeed())

	waitForFirstConsumer := storagev1.VolumeBindingWaitForFirstConsumer
	reclaimDelete := corev1.PersistentVolumeReclaimDelete
	expect := &storagev1.StorageClass{
		VolumeBindingMode:    &waitForFirstConsumer,
		AllowVolumeExpansion: pointer.BoolPtr(true),
		ReclaimPolicy:        &reclaimDelete,
		Provisioner:          "kubernetes.io/gce-pd",
		Parameters: map[string]string{
			"type": "pd-ssd",
		},
		MountOptions: []string{"nodelalloc,noatime"},
	}
	sc.ObjectMeta = expect.ObjectMeta
	sc.TypeMeta = expect.TypeMeta
	g.Expect(expect).Should(Equal(sc))

	// check local storage
	sc = new(storagev1.StorageClass)
	err = client.Get(ctx, types.NamespacedName{Name: "local-storage"}, sc)
	g.Expect(err).Should(Succeed())
	expect = &storagev1.StorageClass{
		Provisioner:       "kubernetes.io/no-provisioner",
		VolumeBindingMode: &waitForFirstConsumer,
		MountOptions:      []string{"nodelalloc,noatime"},
	}
	sc.ObjectMeta = expect.ObjectMeta
	sc.TypeMeta = expect.TypeMeta
	g.Expect(expect).Should(Equal(sc))

	// just check duplicate call will not cause error
	err = ensureStorageClasses(ctx, client, "gcp", "v1.18")
	g.Expect(err).Should(Succeed())
}