package scrape

import (
	"testing"

	tidbClusterV1alpha1 "github.com/pingcap/tidb-operator/pkg/apis/pingcap/v1alpha1"
	"github.com/stretchr/testify/assert"
)

func Test_getComponentReplicas(t *testing.T) {
	tests := []struct {
		name               string
		tc                 *tidbClusterV1alpha1.TidbCluster
		wantTidbReplica    int32
		wantTikvReplica    int32
		wantTiflashReplica int32
	}{
		{
			name:               "nil tc",
			tc:                 nil,
			wantTidbReplica:    0,
			wantTikvReplica:    0,
			wantTiflashReplica: 0,
		},
		{
			name: "1 tidb, 3 tikv",
			tc: &tidbClusterV1alpha1.TidbCluster{
				Spec: tidbClusterV1alpha1.TidbClusterSpec{
					TiDB: &tidbClusterV1alpha1.TiDBSpec{
						Replicas: 1,
					},
					TiKV: &tidbClusterV1alpha1.TiKVSpec{
						Replicas: 3,
					},
				},
			},
			wantTidbReplica:    1,
			wantTikvReplica:    3,
			wantTiflashReplica: 0,
		},
		{
			name: "2 tidb, 3 tikv, 1 tiflash",
			tc: &tidbClusterV1alpha1.TidbCluster{
				Spec: tidbClusterV1alpha1.TidbClusterSpec{
					TiDB: &tidbClusterV1alpha1.TiDBSpec{
						Replicas: 2,
					},
					TiKV: &tidbClusterV1alpha1.TiKVSpec{
						Replicas: 3,
					},
					TiFlash: &tidbClusterV1alpha1.TiFlashSpec{
						Replicas: 1,
					},
				},
			},
			wantTidbReplica:    2,
			wantTikvReplica:    3,
			wantTiflashReplica: 1,
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			gotTidbReplica, gotTikvReplica, gotTiflashReplica := getComponentReplicas(tt.tc)
			assert.Equal(t, tt.wantTidbReplica, gotTidbReplica)
			assert.Equal(t, tt.wantTikvReplica, gotTikvReplica)
			assert.Equal(t, tt.wantTiflashReplica, gotTiflashReplica)
		})
	}
}