Skip to content

Commit d0c0509

Browse files
lhy1024lhy1024
committed
fix some tests
Signed-off-by: lhy1024 <[email protected]>
1 parent 676cfc5 commit d0c0509

File tree

5 files changed

+300
-95
lines changed

5 files changed

+300
-95
lines changed

‎pkg/schedule/checker/affinity_checker_test.go‎

Lines changed: 60 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -612,7 +612,7 @@ func TestAffinityCheckerMergePath(t *testing.T) {
612612
defer cancel()
613613

614614
opt := mockconfig.NewTestOptions()
615-
opt.SetMaxAffinityMergeRegionSize(20)
615+
opt.SetMaxAffinityMergeRegionSize(60) // Larger than default MaxMergeRegionSize
616616
tc := mockcluster.NewCluster(ctx, opt)
617617
tc.AddRegionStore(1, 100)
618618
tc.AddRegionStore(2, 100)
@@ -659,7 +659,7 @@ func TestAffinityMergeCheckNoTarget(t *testing.T) {
659659
defer cancel()
660660

661661
opt := mockconfig.NewTestOptions()
662-
opt.SetMaxAffinityMergeRegionSize(20)
662+
opt.SetMaxAffinityMergeRegionSize(60) // Larger than default MaxMergeRegionSize
663663
tc := mockcluster.NewCluster(ctx, opt)
664664
tc.AddRegionStore(1, 100)
665665
tc.AddRegionStore(2, 100)
@@ -696,7 +696,7 @@ func TestAffinityMergeCheckDifferentGroups(t *testing.T) {
696696
defer cancel()
697697

698698
opt := mockconfig.NewTestOptions()
699-
opt.SetMaxAffinityMergeRegionSize(20)
699+
opt.SetMaxAffinityMergeRegionSize(60) // Larger than default MaxMergeRegionSize
700700
tc := mockcluster.NewCluster(ctx, opt)
701701
tc.AddRegionStore(1, 100)
702702
tc.AddRegionStore(2, 100)
@@ -757,7 +757,7 @@ func TestAffinityMergeCheckRegionTooLarge(t *testing.T) {
757757
defer cancel()
758758

759759
opt := mockconfig.NewTestOptions()
760-
opt.SetMaxAffinityMergeRegionSize(20)
760+
opt.SetMaxAffinityMergeRegionSize(60) // Larger than default MaxMergeRegionSize
761761
tc := mockcluster.NewCluster(ctx, opt)
762762
tc.AddRegionStore(1, 100)
763763
tc.AddRegionStore(2, 100)
@@ -770,8 +770,8 @@ func TestAffinityMergeCheckRegionTooLarge(t *testing.T) {
770770
region2 := tc.GetRegion(2)
771771

772772
region1 = region1.Clone(
773-
core.SetApproximateSize(30), // Too large to merge
774-
core.SetApproximateKeys(30000),
773+
core.SetApproximateSize(70), // Too large to merge under 60 MB limit
774+
core.SetApproximateKeys(700000),
775775
core.WithStartKey([]byte("a")),
776776
core.WithEndKey([]byte("b")),
777777
)
@@ -1012,7 +1012,7 @@ func TestAffinityMergeCheckTargetTooBig(t *testing.T) {
10121012
defer cancel()
10131013

10141014
opt := mockconfig.NewTestOptions()
1015-
opt.SetMaxAffinityMergeRegionSize(20) // Max size 20, Max keys 200000
1015+
opt.SetMaxAffinityMergeRegionSize(60) // Larger than default MaxMergeRegionSize
10161016
tc := mockcluster.NewCluster(ctx, opt)
10171017
tc.AddRegionStore(1, 100)
10181018
tc.AddRegionStore(2, 100)
@@ -1022,13 +1022,63 @@ func TestAffinityMergeCheckTargetTooBig(t *testing.T) {
10221022
tc.AddLeaderRegion(1, 1, 2, 3)
10231023
tc.AddLeaderRegion(2, 1, 2, 3)
10241024
region1 := tc.GetRegion(1).Clone(
1025-
core.SetApproximateSize(15), // 15 size (Source)
1025+
core.SetApproximateSize(45), // Source fits under 60 MB limit
1026+
core.SetApproximateKeys(450000),
1027+
core.WithStartKey([]byte("a")),
1028+
core.WithEndKey([]byte("b")),
1029+
)
1030+
region2 := tc.GetRegion(2).Clone(
1031+
core.SetApproximateSize(25), // Total: 70 > 60
1032+
core.SetApproximateKeys(250000),
1033+
core.WithStartKey([]byte("b")),
1034+
core.WithEndKey([]byte("c")),
1035+
)
1036+
1037+
tc.PutRegion(region1)
1038+
tc.PutRegion(region2)
1039+
1040+
affinityManager := tc.GetAffinityManager()
1041+
checker := NewAffinityChecker(tc, opt)
1042+
1043+
group := &affinity.Group{
1044+
ID: "test_group",
1045+
LeaderStoreID: 1,
1046+
VoterStoreIDs: []uint64{1, 2, 3},
1047+
}
1048+
err := createAffinityGroupForTest(affinityManager, group, []byte(""), []byte(""))
1049+
re.NoError(err)
1050+
1051+
// MergeCheck should return nil because the combined size (70) exceeds the limit (60)
1052+
groupState, _ := affinityManager.GetRegionAffinityGroupState(region1)
1053+
re.NotNil(groupState)
1054+
ops := checker.MergeCheck(region1, groupState)
1055+
re.Nil(ops, "Merged size exceeds MaxAffinityMergeRegionSize")
1056+
}
1057+
1058+
// TestAffinityMergeCheckTargetRespectsLargerGlobalLimit verifies that the merge
1059+
// limit follows the larger value between MaxMergeRegionSize and MaxAffinityMergeRegionSize.
1060+
func TestAffinityMergeCheckTargetRespectsLargerGlobalLimit(t *testing.T) {
1061+
re := require.New(t)
1062+
ctx, cancel := context.WithCancel(context.Background())
1063+
defer cancel()
1064+
1065+
opt := mockconfig.NewTestOptions()
1066+
opt.SetMaxAffinityMergeRegionSize(20) // Smaller than the default MaxMergeRegionSize (54)
1067+
tc := mockcluster.NewCluster(ctx, opt)
1068+
tc.AddRegionStore(1, 100)
1069+
tc.AddRegionStore(2, 100)
1070+
tc.AddRegionStore(3, 100)
1071+
1072+
tc.AddLeaderRegion(1, 1, 2, 3)
1073+
tc.AddLeaderRegion(2, 1, 2, 3)
1074+
region1 := tc.GetRegion(1).Clone(
1075+
core.SetApproximateSize(15),
10261076
core.SetApproximateKeys(150000),
10271077
core.WithStartKey([]byte("a")),
10281078
core.WithEndKey([]byte("b")),
10291079
)
10301080
region2 := tc.GetRegion(2).Clone(
1031-
core.SetApproximateSize(6), // 6 size (Target). Total: 21 > 20
1081+
core.SetApproximateSize(6), // Total size: 21, allowed by global limit 54
10321082
core.SetApproximateKeys(60000),
10331083
core.WithStartKey([]byte("b")),
10341084
core.WithEndKey([]byte("c")),
@@ -1048,11 +1098,10 @@ func TestAffinityMergeCheckTargetTooBig(t *testing.T) {
10481098
err := createAffinityGroupForTest(affinityManager, group, []byte(""), []byte(""))
10491099
re.NoError(err)
10501100

1051-
// MergeCheck should return nil because the combined size (21) exceeds the limit (20)
10521101
groupState, _ := affinityManager.GetRegionAffinityGroupState(region1)
10531102
re.NotNil(groupState)
10541103
ops := checker.MergeCheck(region1, groupState)
1055-
re.Nil(ops, "Merged size exceeds MaxAffinityMergeRegionSize")
1104+
re.NotNil(ops, "Merge should be allowed because the effective limit follows the larger global max")
10561105
}
10571106

10581107
// TestAffinityMergeCheckAdjacentUnhealthy tests that merging is blocked if the adjacent region is unhealthy.

‎pkg/utils/apiutil/serverapi/middleware.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ func (h *redirector) matchMicroserviceRedirectRules(r *http.Request) (bool, stri
126126
}
127127

128128
// MatchMicroserviceRedirect checks rules, rewrites path in-place, and returns (matched, targetAddr).
129-
// If matched but no primary is available, it returns matched=false to let caller handle locally.
129+
// If matched but no primary is available, it returns matched=true with empty addr so caller can handle the redirect error.
130130
func MatchMicroserviceRedirect(
131131
r *http.Request,
132132
rules []RedirectRule,
@@ -160,7 +160,7 @@ func MatchMicroserviceRedirect(
160160
if !ok || addr == "" {
161161
log.Warn("failed to get the service primary addr when trying to match redirect rules",
162162
zap.String("path", r.URL.Path))
163-
continue
163+
return true, ""
164164
}
165165
// If the URL contains escaped characters, use RawPath instead of Path
166166
origin := r.URL.Path

‎pkg/utils/apiutil/serverapi/middleware_test.go‎

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
// Copyright 2025 TiKV Project Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package serverapi
16+
17+
import (
18+
"context"
19+
"net/http"
20+
"testing"
21+
22+
"github.com/stretchr/testify/require"
23+
"go.uber.org/goleak"
24+
25+
"github.com/tikv/pd/pkg/mcs/utils/constant"
26+
"github.com/tikv/pd/pkg/utils/apiutil"
27+
"github.com/tikv/pd/pkg/utils/testutil"
28+
)
29+
30+
func TestMain(m *testing.M) {
31+
goleak.VerifyTestMain(m, testutil.LeakOptions...)
32+
}
33+
34+
// Ensure when a rule matches but no primary address is available, we still mark it matched
35+
// so the caller can return ErrRedirect (expected by API behavior).
36+
func TestMatchMicroserviceRedirectNoPrimary(t *testing.T) {
37+
re := require.New(t)
38+
req, err := http.NewRequest(http.MethodPost, "/pd/api/v1/admin/reset-ts", http.NoBody)
39+
re.NoError(err)
40+
41+
matched, addr := MatchMicroserviceRedirect(
42+
req,
43+
[]RedirectRule{{
44+
MatchPath: "/pd/api/v1/admin/reset-ts",
45+
TargetPath: "/tso/api/v1/admin/reset-ts",
46+
TargetServiceName: constant.TSOServiceName,
47+
MatchMethods: []string{http.MethodPost},
48+
}},
49+
true, /* isKeyspaceGroupEnabled */
50+
func(string) bool { return true }, /* isServiceIndependent */
51+
func(_ context.Context, _ string) (string, bool) { return "", false }, /* getPrimary */
52+
)
53+
re.True(matched)
54+
re.Empty(addr)
55+
}
56+
57+
func TestMatchMicroserviceRedirectForbiddenHeader(t *testing.T) {
58+
re := require.New(t)
59+
req, err := http.NewRequest(http.MethodGet, "/pd/api/v1/foo", http.NoBody)
60+
re.NoError(err)
61+
req.Header.Set(apiutil.XForbiddenForwardToMicroserviceHeader, "true")
62+
63+
matched, addr := MatchMicroserviceRedirect(
64+
req,
65+
[]RedirectRule{{
66+
MatchPath: "/pd/api/v1/foo",
67+
TargetPath: "/target",
68+
TargetServiceName: constant.SchedulingServiceName,
69+
MatchMethods: []string{http.MethodGet},
70+
}},
71+
true,
72+
func(string) bool { return true },
73+
func(_ context.Context, _ string) (string, bool) { return "addr", true },
74+
)
75+
76+
re.False(matched)
77+
re.Empty(addr)
78+
}
79+
80+
func TestMatchMicroserviceRedirectKeyspaceDisabled(t *testing.T) {
81+
re := require.New(t)
82+
req, err := http.NewRequest(http.MethodGet, "/pd/api/v1/foo", http.NoBody)
83+
re.NoError(err)
84+
85+
matched, addr := MatchMicroserviceRedirect(
86+
req,
87+
[]RedirectRule{{
88+
MatchPath: "/pd/api/v1/foo",
89+
TargetPath: "/target",
90+
TargetServiceName: constant.SchedulingServiceName,
91+
MatchMethods: []string{http.MethodGet},
92+
}},
93+
false, // keyspace group disabled
94+
func(string) bool { return true },
95+
func(_ context.Context, _ string) (string, bool) { return "addr", true },
96+
)
97+
98+
re.False(matched)
99+
re.Empty(addr)
100+
}
101+
102+
func TestMatchMicroserviceRedirectSchedulingNotIndependent(t *testing.T) {
103+
re := require.New(t)
104+
req, err := http.NewRequest(http.MethodGet, "/pd/api/v1/foo", http.NoBody)
105+
re.NoError(err)
106+
107+
matched, addr := MatchMicroserviceRedirect(
108+
req,
109+
[]RedirectRule{{
110+
MatchPath: "/pd/api/v1/foo",
111+
TargetPath: "/target",
112+
TargetServiceName: constant.SchedulingServiceName,
113+
MatchMethods: []string{http.MethodGet},
114+
}},
115+
true,
116+
func(name string) bool { return name != constant.SchedulingServiceName },
117+
func(_ context.Context, _ string) (string, bool) { return "addr", true },
118+
)
119+
120+
re.False(matched)
121+
re.Empty(addr)
122+
}
123+
124+
func TestMatchMicroserviceRedirectFilterBlocks(t *testing.T) {
125+
re := require.New(t)
126+
req, err := http.NewRequest(http.MethodGet, "/pd/api/v1/foo", http.NoBody)
127+
re.NoError(err)
128+
129+
matched, addr := MatchMicroserviceRedirect(
130+
req,
131+
[]RedirectRule{{
132+
MatchPath: "/pd/api/v1/foo",
133+
TargetPath: "/target",
134+
TargetServiceName: constant.SchedulingServiceName,
135+
MatchMethods: []string{http.MethodGet},
136+
Filter: func(_ *http.Request) bool {
137+
return false
138+
},
139+
}},
140+
true,
141+
func(string) bool { return true },
142+
func(_ context.Context, _ string) (string, bool) { return "addr", true },
143+
)
144+
145+
re.False(matched)
146+
re.Empty(addr)
147+
}

0 commit comments

Comments
 (0)