chore: add UT for PD/TiKV/TiDB/TiFlash APIs (#6012)

pkg/tidbapi/v1/client_test.go

@@ -0,0 +1,87 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tidbapi
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTiDBClient_GetHealth(t *testing.T) {
+	jsonStr := `{"connections":0,"version":"8.0.11-TiDB-v8.1.0","git_hash":"945d07c5d5c7a1ae212f6013adfb187f2de24b23"}`
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/status", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte(jsonStr))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	client := NewTiDBClient(server.URL, 5*time.Second, nil)
+	health, err := client.GetHealth(context.Background())
+	require.NoError(t, err)
+	assert.True(t, health)
+}
+
+func TestTiDBClient_GetInfo(t *testing.T) {
+	jsonStr := `
+{
+ "is_owner": true,
+ "max_procs": 10,
+ "gogc": 500,
+ "version": "8.0.11-TiDB-v8.1.0",
+ "git_hash": "945d07c5d5c7a1ae212f6013adfb187f2de24b23",
+ "ddl_id": "fe4de332-a1c5-46ba-a1d0-762c716345d3",
+ "ip": "basic-tidb-9lbgl4.basic-tidb-peer.default.svc",
+ "listening_port": 4000,
+ "status_port": 10080,
+ "lease": "45s",
+ "binlog_status": "Off",
+ "start_timestamp": 1735095910,
+ "labels": {},
+ "server_id": 85
+}`
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/info", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte(jsonStr))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	client := NewTiDBClient(server.URL, 5*time.Second, nil)
+	info, err := client.GetInfo(context.Background())
+	require.NoError(t, err)
+	assert.True(t, info.IsOwner)
+}
+
+func TestTiDBClient_SetServerLabels(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/labels", r.URL.Path)
+		assert.Equal(t, http.MethodPost, r.Method)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer server.Close()
+
+	client := NewTiDBClient(server.URL, 5*time.Second, nil)
+	err := client.SetServerLabels(context.Background(), map[string]string{"region": "us-west-1"})
+	require.NoError(t, err)
+}

pkg/tidbapi/v1/control.go

@@ -20,8 +20,6 @@ import (
 	"fmt"
 	"time"
 
-	corelisterv1 "k8s.io/client-go/listers/core/v1"
-
 	"github.com/pingcap/tidb-operator/apis/core/v1alpha1"
 	"github.com/pingcap/tidb-operator/pkg/client"
 	tlsutil "github.com/pingcap/tidb-operator/pkg/utils/tls"
@@ -41,12 +39,11 @@ type TiDBControlInterface interface {
 
 // defaultTiDBControl is the default implementation of TiDBControlInterface.
 type defaultTiDBControl struct {
-	secretLister corelisterv1.SecretLister
 }
 
 // NewDefaultTiDBControl returns a defaultTiDBControl instance.
-func NewDefaultTiDBControl(secretLister corelisterv1.SecretLister) TiDBControlInterface {
-	return &defaultTiDBControl{secretLister: secretLister}
+func NewDefaultTiDBControl() TiDBControlInterface {
+	return &defaultTiDBControl{}
 }
 
 // GetTiDBPodClient provides TiDBClient of a TiDB pod.

pkg/tidbapi/v1/control_test.go

@@ -0,0 +1,44 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tidbapi
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTiDDBControl(t *testing.T) {
+	jsonStr := `{"connections":0,"version":"8.0.11-TiDB-v8.1.0","git_hash":"945d07c5d5c7a1ae212f6013adfb187f2de24b23"}`
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/status", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte(jsonStr))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	// only test non-tls
+	client, err := NewDefaultTiDBControl().GetTiDBPodClient(context.Background(), nil, "", "", "", "", false)
+	require.NoError(t, err)
+	client.(*tidbClient).url = server.URL // replace the url with the test server
+	health, err := client.GetHealth(context.Background())
+	require.NoError(t, err)
+	assert.True(t, health)
+}

pkg/tiflashapi/v1/client_test.go

@@ -0,0 +1,41 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tiflashapi
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTiFlash_GetStoreStatus(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/tiflash/store-status", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte("Running"))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	client := NewTiFlashClient(server.URL, 5*time.Second, nil, false)
+	status, err := client.GetStoreStatus(context.Background())
+	require.NoError(t, err)
+	assert.Equal(t, Running, status)
+}

pkg/tiflashapi/v1/control.go

@@ -20,8 +20,6 @@ import (
 	"fmt"
 	"time"
 
-	corelisterv1 "k8s.io/client-go/listers/core/v1"
-
 	"github.com/pingcap/tidb-operator/apis/core/v1alpha1"
 	"github.com/pingcap/tidb-operator/pkg/client"
 	tlsutil "github.com/pingcap/tidb-operator/pkg/utils/tls"
@@ -41,12 +39,11 @@ type TiFlashControlInterface interface {
 
 // defaultTiFlashControl is the default implementation of TiFlashControlInterface.
 type defaultTiFlashControl struct {
-	secretLister corelisterv1.SecretLister
 }
 
 // NewDefaultTiFlashControl returns a defaultTiFlashControl instance
-func NewDefaultTiFlashControl(secretLister corelisterv1.SecretLister) TiFlashControlInterface {
-	return &defaultTiFlashControl{secretLister: secretLister}
+func NewDefaultTiFlashControl() TiFlashControlInterface {
+	return &defaultTiFlashControl{}
 }
 
 // GetTiFlashPodClient provides TiFlashClient of a TiFlash pod.

pkg/tiflashapi/v1/control_test.go

@@ -0,0 +1,43 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tiflashapi
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTiFlashControl(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/tiflash/store-status", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte("Ready"))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	// only test non-tls
+	client, err := NewDefaultTiFlashControl().GetTiFlashPodClient(context.Background(), nil, "", "", "", false)
+	require.NoError(t, err)
+	client.(*tiflashClient).url = server.URL
+	status, err := client.GetStoreStatus(context.Background())
+	require.NoError(t, err)
+	assert.Equal(t, Ready, status)
+}

pkg/tikvapi/v1/client_test.go

@@ -0,0 +1,47 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tikvapi
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTiKVClient_GetLeaderCount(t *testing.T) {
+	jsonStr := `
+# HELP tikv_raftstore_region_count Number of regions collected in region_collector
+# TYPE tikv_raftstore_region_count gauge
+tikv_raftstore_region_count{type="buckets"} 50
+tikv_raftstore_region_count{type="leader"} 20
+tikv_raftstore_region_count{type="region"} 50
+`
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/metrics", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte(jsonStr))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	client := NewTiKVClient(server.URL, 5*time.Second, nil, false)
+	count, err := client.GetLeaderCount()
+	require.NoError(t, err)
+	assert.Equal(t, 20, count)
+}

pkg/tikvapi/v1/control.go

@@ -20,8 +20,6 @@ import (
 	"fmt"
 	"time"
 
-	corelisterv1 "k8s.io/client-go/listers/core/v1"
-
 	"github.com/pingcap/tidb-operator/apis/core/v1alpha1"
 	"github.com/pingcap/tidb-operator/pkg/client"
 	tlsutil "github.com/pingcap/tidb-operator/pkg/utils/tls"
@@ -41,12 +39,11 @@ type TiKVControlInterface interface {
 
 // defaultTiKVControl is the default implementation of TiKVControlInterface.
 type defaultTiKVControl struct {
-	secretLister corelisterv1.SecretLister
 }
 
 // NewDefaultTiKVControl returns a defaultTiKVControl instance.
-func NewDefaultTiKVControl(secretLister corelisterv1.SecretLister) TiKVControlInterface {
-	return &defaultTiKVControl{secretLister: secretLister}
+func NewDefaultTiKVControl() TiKVControlInterface {
+	return &defaultTiKVControl{}
 }
 
 // GetTiKVPodClient provides TiKVClient of a TiKV pod.

pkg/tikvapi/v1/control_test.go

@@ -0,0 +1,50 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tikvapi
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestTiKVControl(t *testing.T) {
+	jsonStr := `
+# HELP tikv_raftstore_region_count Number of regions collected in region_collector
+# TYPE tikv_raftstore_region_count gauge
+tikv_raftstore_region_count{type="buckets"} 50
+tikv_raftstore_region_count{type="leader"} 20
+tikv_raftstore_region_count{type="region"} 50
+`
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		assert.Equal(t, "/metrics", r.URL.Path)
+		assert.Equal(t, http.MethodGet, r.Method)
+		_, err := w.Write([]byte(jsonStr))
+		assert.NoError(t, err)
+	}))
+	defer server.Close()
+
+	// only test non-tls
+	client, err := NewDefaultTiKVControl().GetTiKVPodClient(context.Background(), nil, "", "", "", "", false)
+	require.NoError(t, err)
+	client.(*tikvClient).url = server.URL
+	count, err := client.GetLeaderCount()
+	require.NoError(t, err)
+	assert.Equal(t, 20, count)
+}