package models

import (
	"testing"
)

func TestGetStorageTierByDiskType(t *testing.T) {
	tests := []struct {
		name     string
		diskType string
		want     string
	}{
		{
			name:     "io2 should return plus",
			diskType: "io2",
			want:     NextGenStorageTierPlus,
		},
		{
			name:     "gp3 should return standard",
			diskType: "gp3",
			want:     NextGenStorageTierStandard,
		},
		{
			name:     "unknown disk type should return standard",
			diskType: "unknown",
			want:     NextGenStorageTierStandard,
		},
		{
			name:     "gp2 should return standard",
			diskType: "gp2",
			want:     NextGenStorageTierStandard,
		},
		{
			name:     "empty disk type should return standard",
			diskType: "",
			want:     NextGenStorageTierStandard,
		},
		{
			name:     "io1 should return standard",
			diskType: "io1",
			want:     NextGenStorageTierStandard,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := GetStorageTierByDiskType(tt.diskType)
			if got != tt.want {
				t.Errorf("GetStorageTierByDiskType(%q) = %q, want %q", tt.diskType, got, tt.want)
			}
		})
	}
}

func TestGetStorageLayoutByComponent(t *testing.T) {
	tests := []struct {
		name      string
		component string
		want      string
	}{
		{
			name:      "tikv should return row",
			component: "tikv",
			want:      NextGenStorageLayoutRow,
		},
		{
			name:      "tikv-log should return row",
			component: "tikv-log",
			want:      NextGenStorageLayoutRow,
		},
		{
			name:      "tiflash should return column",
			component: "tiflash",
			want:      NextGenStorageLayoutColumn,
		},
		{
			name:      "tiflash-compute should return column",
			component: "tiflash-compute",
			want:      NextGenStorageLayoutColumn,
		},
		{
			name:      "tiflash-write should return column",
			component: "tiflash-write",
			want:      NextGenStorageLayoutColumn,
		},
		{
			name:      "unknown component should return column",
			component: "unknown",
			want:      NextGenStorageLayoutColumn,
		},
		{
			name:      "empty component should return column",
			component: "",
			want:      NextGenStorageLayoutColumn,
		},
		{
			name:      "tidb should return column",
			component: "tidb",
			want:      NextGenStorageLayoutColumn,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			got := GetStorageLayoutByComponent(tt.component)
			if got != tt.want {
				t.Errorf("GetStorageLayoutByComponent(%q) = %q, want %q", tt.component, got, tt.want)
			}
		})
	}
}