Skip to content

Commit 8015da8

Browse files
tallclairtallclair
committed
Refactor test agent to interface
1 parent e0e1514 commit 8015da8

12 files changed

+209
-155
lines changed

‎tests/agent_disconnect_test.go

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -57,15 +57,13 @@ func TestProxy_Agent_Disconnect_HTTP_Persistent_Connection(t *testing.T) {
5757
server := httptest.NewServer(newEchoServer("hello"))
5858
defer server.Close()
5959

60-
stopCh := make(chan struct{})
61-
6260
proxy, cleanup, err := tc.proxyServerFunction()
6361
if err != nil {
6462
t.Fatal(err)
6563
}
6664
defer cleanup()
6765

68-
runAgent(proxy.agent, stopCh)
66+
a := runAgent(t, proxy.agent)
6967
waitForConnectedAgentCount(t, 1, proxy.server)
7068

7169
// run test client
@@ -80,7 +78,7 @@ func TestProxy_Agent_Disconnect_HTTP_Persistent_Connection(t *testing.T) {
8078
if err != nil {
8179
t.Errorf("expected no error on proxy request, got %v", err)
8280
}
83-
close(stopCh)
81+
a.Stop()
8482

8583
// Wait for the agent to disconnect
8684
waitForConnectedAgentCount(t, 0, proxy.server)
@@ -121,16 +119,14 @@ func TestProxy_Agent_Reconnect(t *testing.T) {
121119
server := httptest.NewServer(newEchoServer("hello"))
122120
defer server.Close()
123121

124-
stopCh := make(chan struct{})
125-
126122
proxy, cleanup, err := tc.proxyServerFunction()
127123
if err != nil {
128124
t.Fatal(err)
129125
}
130126
defer cleanup()
131127

132-
cs1 := runAgent(proxy.agent, stopCh)
133-
waitForConnectedServerCount(t, 1, cs1)
128+
ai1 := runAgent(t, proxy.agent)
129+
waitForConnectedServerCount(t, 1, ai1)
134130

135131
// run test client
136132

@@ -143,16 +139,15 @@ func TestProxy_Agent_Reconnect(t *testing.T) {
143139
if err != nil {
144140
t.Errorf("expected no error on proxy request, got %v", err)
145141
}
146-
close(stopCh)
142+
ai1.Stop()
147143

148144
// Wait for the agent to disconnect
149145
waitForConnectedAgentCount(t, 0, proxy.server)
150146

151147
// Reconnect agent
152-
stopCh2 := make(chan struct{})
153-
defer close(stopCh2)
154-
cs2 := runAgent(proxy.agent, stopCh2)
155-
waitForConnectedServerCount(t, 1, cs2)
148+
ai2 := runAgent(t, proxy.agent)
149+
defer ai2.Stop()
150+
waitForConnectedServerCount(t, 1, ai2)
156151

157152
// Proxy requests should work again after agent reconnects
158153
c2, err := tc.clientFunction(ctx, proxy.front, server.URL)

‎tests/benchmarks_test.go

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -39,17 +39,14 @@ func BenchmarkLargeResponse_GRPC(b *testing.B) {
3939
server := httptest.NewServer(newSizedServer(length, chunks))
4040
defer server.Close()
4141

42-
stopCh := make(chan struct{})
43-
defer close(stopCh)
44-
4542
proxy, cleanup, err := runGRPCProxyServer()
4643
if err != nil {
4744
b.Fatal(err)
4845
}
4946
defer cleanup()
5047

51-
clientset := runAgent(proxy.agent, stopCh)
52-
waitForConnectedServerCount(b, 1, clientset)
48+
a := runAgent(b, proxy.agent)
49+
waitForConnectedServerCount(b, 1, a)
5350

5451
req, err := http.NewRequest("GET", server.URL, nil)
5552
if err != nil {
@@ -113,17 +110,14 @@ func BenchmarkLargeRequest_GRPC(b *testing.B) {
113110
}))
114111
defer server.Close()
115112

116-
stopCh := make(chan struct{})
117-
defer close(stopCh)
118-
119113
proxy, cleanup, err := runGRPCProxyServer()
120114
if err != nil {
121115
b.Fatal(err)
122116
}
123117
defer cleanup()
124118

125-
clientset := runAgent(proxy.agent, stopCh)
126-
waitForConnectedServerCount(b, 1, clientset)
119+
a := runAgent(b, proxy.agent)
120+
waitForConnectedServerCount(b, 1, a)
127121

128122
bodyBytes := make([]byte, length)
129123
body := bytes.NewReader(bodyBytes)

‎tests/concurrent_client_request_test.go

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -140,13 +140,13 @@ func TestConcurrentClientRequest(t *testing.T) {
140140
defer cleanup()
141141
ps.BackendManagers = []server.BackendManager{newSingleTimeGetter(server.NewDefaultBackendManager())}
142142

143-
stopCh := make(chan struct{})
144-
defer close(stopCh)
145143
// Run two agents
146-
cs1 := runAgentWithID("a", proxy.agent, stopCh)
147-
cs2 := runAgentWithID("b", proxy.agent, stopCh)
148-
waitForConnectedServerCount(t, 1, cs1)
149-
waitForConnectedServerCount(t, 1, cs2)
144+
ai1 := runAgent(t, proxy.agent)
145+
ai2 := runAgent(t, proxy.agent)
146+
defer ai1.Stop()
147+
defer ai2.Stop()
148+
waitForConnectedServerCount(t, 1, ai1)
149+
waitForConnectedServerCount(t, 1, ai2)
150150

151151
client1 := getTestClient(proxy.front, t)
152152
client2 := getTestClient(proxy.front, t)

‎tests/concurrent_test.go

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -35,17 +35,15 @@ func TestProxy_ConcurrencyGRPC(t *testing.T) {
3535
server := httptest.NewServer(newSizedServer(length, chunks))
3636
defer server.Close()
3737

38-
stopCh := make(chan struct{})
39-
defer close(stopCh)
40-
4138
proxy, cleanup, err := runGRPCProxyServer()
4239
if err != nil {
4340
t.Fatal(err)
4441
}
4542
defer cleanup()
4643

47-
clientset := runAgent(proxy.agent, stopCh)
48-
waitForConnectedServerCount(t, 1, clientset)
44+
a := runAgent(t, proxy.agent)
45+
defer a.Stop()
46+
waitForConnectedServerCount(t, 1, a)
4947

5048
var wg sync.WaitGroup
5149
verify := func() {
@@ -97,17 +95,15 @@ func TestProxy_ConcurrencyHTTP(t *testing.T) {
9795
server := httptest.NewServer(newSizedServer(length, chunks))
9896
defer server.Close()
9997

100-
stopCh := make(chan struct{})
101-
defer close(stopCh)
102-
10398
proxy, cleanup, err := runHTTPConnProxyServer()
10499
if err != nil {
105100
t.Fatal(err)
106101
}
107102
defer cleanup()
108103

109-
clientset := runAgent(proxy.agent, stopCh)
110-
waitForConnectedServerCount(t, 1, clientset)
104+
a := runAgent(t, proxy.agent)
105+
defer a.Stop()
106+
waitForConnectedServerCount(t, 1, a)
111107

112108
// run test clients
113109
var wg sync.WaitGroup
@@ -161,18 +157,15 @@ func TestAgent_MultipleConn(t *testing.T) {
161157
server := httptest.NewServer(waitServer)
162158
defer server.Close()
163159

164-
stopCh := make(chan struct{})
165-
stopCh2 := make(chan struct{})
166-
167160
proxy, cleanup, err := tc.proxyServerFunction()
168161
if err != nil {
169162
t.Fatal(err)
170163
}
171164
defer cleanup()
172165

173-
cs1 := runAgentWithID("multipleAgentConn", proxy.agent, stopCh)
174-
waitForConnectedServerCount(t, 1, cs1)
175-
defer close(stopCh)
166+
ai1 := runAgentWithID(t, "multipleAgentConn", proxy.agent)
167+
defer ai1.Stop()
168+
waitForConnectedServerCount(t, 1, ai1)
176169

177170
// run test client
178171
c, err := tc.clientFunction(ctx, proxy.front, server.URL)
@@ -194,9 +187,9 @@ func TestAgent_MultipleConn(t *testing.T) {
194187
// Running an agent with the same ID simulates a second connection from the same agent.
195188
// This simulates the scenario where a proxy agent established connections with HA proxy server
196189
// and creates multiple connections with the same proxy server
197-
cs2 := runAgentWithID("multipleAgentConn", proxy.agent, stopCh2)
198-
waitForConnectedServerCount(t, 1, cs2)
199-
close(stopCh2)
190+
ai2 := runAgentWithID(t, "multipleAgentConn", proxy.agent)
191+
defer ai2.Stop()
192+
waitForConnectedServerCount(t, 1, ai2)
200193
// Wait for the server to run cleanup routine
201194
waitForConnectedAgentCount(t, 1, proxy.server)
202195
close(waitServer.respondCh)

‎tests/framework/agent.go

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
/*
2+
Copyright 2023 The Kubernetes Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package framework
18+
19+
import (
20+
"time"
21+
22+
"google.golang.org/grpc"
23+
"sigs.k8s.io/apiserver-network-proxy/pkg/agent"
24+
)
25+
26+
type AgentOpts struct {
27+
AgentID string
28+
Addr string
29+
}
30+
31+
type AgentRunner interface {
32+
Start(AgentOpts) (Agent, error)
33+
}
34+
35+
type Agent interface {
36+
GetConnectedServerCount() (int, error)
37+
Ready() bool
38+
Stop() error
39+
}
40+
41+
type InProcessAgentRunner struct{}
42+
43+
func (*InProcessAgentRunner) Start(opts AgentOpts) (Agent, error) {
44+
cc := agent.ClientSetConfig{
45+
Address: opts.Addr,
46+
AgentID: opts.AgentID,
47+
SyncInterval: 100 * time.Millisecond,
48+
ProbeInterval: 100 * time.Millisecond,
49+
DialOptions: []grpc.DialOption{grpc.WithInsecure()},
50+
}
51+
52+
stopCh := make(chan struct{})
53+
client := cc.NewAgentClientSet(stopCh)
54+
client.Serve()
55+
56+
return &inProcessAgent{
57+
client: client,
58+
stopCh: stopCh,
59+
}, nil
60+
}
61+
62+
type inProcessAgent struct {
63+
client *agent.ClientSet
64+
stopCh chan struct{}
65+
}
66+
67+
func (a *inProcessAgent) Stop() error {
68+
close(a.stopCh)
69+
return nil
70+
}
71+
72+
func (a *inProcessAgent) GetConnectedServerCount() (int, error) {
73+
return a.client.HealthyClientsCount(), nil
74+
}
75+
76+
func (a *inProcessAgent) Ready() bool {
77+
return a.client.Ready()
78+
}

‎tests/framework/framework.go

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
/*
2+
Copyright 2023 The Kubernetes Authors.
3+
4+
Licensed under the Apache License, Version 2.0 (the "License");
5+
you may not use this file except in compliance with the License.
6+
You may obtain a copy of the License at
7+
8+
http://www.apache.org/licenses/LICENSE-2.0
9+
10+
Unless required by applicable law or agreed to in writing, software
11+
distributed under the License is distributed on an "AS IS" BASIS,
12+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
See the License for the specific language governing permissions and
14+
limitations under the License.
15+
*/
16+
17+
package framework
18+
19+
type Framework struct {
20+
AgentRunner AgentRunner
21+
}

‎tests/ha_proxy_server_test.go

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -148,11 +148,9 @@ func TestBasicHAProxyServer_GRPC(t *testing.T) {
148148
}
149149
lbAddr := lb.serve(stopCh)
150150

151-
clientset := runAgent(lbAddr, stopCh)
152-
waitForConnectedServerCount(t, 3, clientset)
153-
if cc := clientset.ClientsCount(); cc != 3 {
154-
t.Fatalf("Expected 3 clients, got %d", cc)
155-
}
151+
ai := runAgent(t, lbAddr)
152+
defer ai.Stop()
153+
waitForConnectedServerCount(t, 3, ai)
156154

157155
// run test client
158156
testProxyServer(t, proxy[0].front, server.URL)
@@ -166,7 +164,7 @@ func TestBasicHAProxyServer_GRPC(t *testing.T) {
166164
cleanups[0]()
167165

168166
// give the agent some time to detect the disconnection
169-
waitForConnectedServerCount(t, 2, clientset)
167+
waitForConnectedServerCount(t, 2, ai)
170168

171169
proxy4, _, cleanup4, err := runGRPCProxyServerWithServerCount(haServerCount)
172170
if err != nil {
@@ -180,10 +178,7 @@ func TestBasicHAProxyServer_GRPC(t *testing.T) {
180178
}()
181179

182180
// wait for the new server to be connected.
183-
waitForConnectedServerCount(t, 3, clientset)
184-
if cc := clientset.ClientsCount(); cc != 3 && cc != 4 {
185-
t.Fatalf("Expected 3 or 4 clients, got %d", cc)
186-
}
181+
waitForConnectedServerCount(t, 3, ai)
187182

188183
// run test client
189184
testProxyServer(t, proxy[1].front, server.URL)

‎tests/main_test.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,13 @@ import (
2424
"k8s.io/klog/v2"
2525

2626
metricsclient "sigs.k8s.io/apiserver-network-proxy/konnectivity-client/pkg/client/metrics"
27+
"sigs.k8s.io/apiserver-network-proxy/tests/framework"
2728
)
2829

30+
var Framework = framework.Framework{
31+
AgentRunner: &framework.InProcessAgentRunner{},
32+
}
33+
2934
func TestMain(m *testing.M) {
3035
fs := flag.NewFlagSet("mock-flags", flag.PanicOnError)
3136
klog.InitFlags(fs)

0 commit comments

Comments
 (0)