Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit10ee938

Browse files
committed
chore: more tests
1 parent55e6436 commit10ee938

File tree

9 files changed

+422
-266
lines changed

9 files changed

+422
-266
lines changed

‎enterprise/x/aibridged/aibridged.go‎

Lines changed: 34 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -68,78 +68,78 @@ func New(ctx context.Context, pool Pooler, rpcDialer Dialer, logger slog.Logger)
6868
}
6969

7070
// Connect establishes a connection to coderd.
71-
func (d*Server)connect() {
72-
deferd.logger.Debug(d.lifecycleCtx,"connect loop exited")
73-
deferd.wg.Done()
71+
func (s*Server)connect() {
72+
defers.logger.Debug(s.lifecycleCtx,"connect loop exited")
73+
defers.wg.Done()
7474

75-
logConnect:=d.logger.With(slog.F("context","aibridged.server")).Debug
75+
logConnect:=s.logger.With(slog.F("context","aibridged.server")).Debug
7676
// An exponential back-off occurs when the connection is failing to dial.
7777
// This is to prevent server spam in case of a coderd outage.
7878
connectLoop:
79-
forretrier:=retry.New(50*time.Millisecond,10*time.Second);retrier.Wait(d.lifecycleCtx); {
79+
forretrier:=retry.New(50*time.Millisecond,10*time.Second);retrier.Wait(s.lifecycleCtx); {
8080
// It's possible for the aibridge daemon to be shut down
8181
// before the wait is complete!
82-
ifd.isShutdown() {
82+
ifs.isShutdown() {
8383
return
8484
}
85-
d.logger.Debug(d.lifecycleCtx,"dialing coderd")
86-
client,err:=d.clientDialer(d.lifecycleCtx)
85+
s.logger.Debug(s.lifecycleCtx,"dialing coderd")
86+
client,err:=s.clientDialer(s.lifecycleCtx)
8787
iferr!=nil {
8888
iferrors.Is(err,context.Canceled) {
8989
return
9090
}
9191
varsdkErr*codersdk.Error
9292
// If something is wrong with our auth, stop trying to connect.
9393
iferrors.As(err,&sdkErr)&&sdkErr.StatusCode()==http.StatusForbidden {
94-
d.logger.Error(d.lifecycleCtx,"not authorized to dial coderd",slog.Error(err))
94+
s.logger.Error(s.lifecycleCtx,"not authorized to dial coderd",slog.Error(err))
9595
return
9696
}
97-
ifd.isShutdown() {
97+
ifs.isShutdown() {
9898
return
9999
}
100-
d.logger.Warn(d.lifecycleCtx,"coderd client failed to dial",slog.Error(err))
100+
s.logger.Warn(s.lifecycleCtx,"coderd client failed to dial",slog.Error(err))
101101
continue
102102
}
103103

104104
// TODO: log this with INFO level when we implement external aibridge daemons.
105-
logConnect(d.lifecycleCtx,"successfully connected to coderd")
105+
logConnect(s.lifecycleCtx,"successfully connected to coderd")
106106
retrier.Reset()
107-
d.initConnectionOnce.Do(func() {
108-
close(d.initConnectionCh)
107+
s.initConnectionOnce.Do(func() {
108+
close(s.initConnectionCh)
109109
})
110110

111111
// Serve the client until we are closed or it disconnects.
112112
for {
113113
select {
114-
case<-d.lifecycleCtx.Done():
114+
case<-s.lifecycleCtx.Done():
115115
client.DRPCConn().Close()
116116
return
117117
case<-client.DRPCConn().Closed():
118-
logConnect(d.lifecycleCtx,"connection to coderd closed")
118+
logConnect(s.lifecycleCtx,"connection to coderd closed")
119119
continue connectLoop
120-
cased.clientCh<-client:
120+
cases.clientCh<-client:
121121
continue
122122
}
123123
}
124124
}
125125
}
126126

127-
func (d*Server)Client() (DRPCClient,error) {
127+
func (s*Server)Client() (DRPCClient,error) {
128128
select {
129-
case<-d.lifecycleCtx.Done():
129+
case<-s.lifecycleCtx.Done():
130130
returnnil,xerrors.New("context closed")
131-
caseclient:=<-d.clientCh:
131+
caseclient:=<-s.clientCh:
132132
returnclient,nil
133133
}
134134
}
135135

136136
// GetRequestHandler retrieves a (possibly reused) [*aibridge.RequestBridge] from the pool, for the given user.
137-
func (d*Server)GetRequestHandler(ctx context.Context,reqRequest) (http.Handler,error) {
138-
ifd.requestBridgePool==nil {
137+
func (s*Server)GetRequestHandler(ctx context.Context,reqRequest) (http.Handler,error) {
138+
ifs.requestBridgePool==nil {
139139
returnnil,xerrors.New("nil requestBridgePool")
140140
}
141141

142-
reqBridge,err:=d.requestBridgePool.Acquire(ctx,req,d.Client)
142+
reqBridge,err:=s.requestBridgePool.Acquire(ctx,req,s.Client,NewMCPProxyFactory(s.logger,s.Client))
143143
iferr!=nil {
144144
returnnil,xerrors.Errorf("acquire request bridge: %w",err)
145145
}
@@ -148,38 +148,38 @@ func (d *Server) GetRequestHandler(ctx context.Context, req Request) (http.Handl
148148
}
149149

150150
// isShutdown returns whether the Server is shutdown or not.
151-
func (d*Server)isShutdown()bool {
151+
func (s*Server)isShutdown()bool {
152152
select {
153-
case<-d.lifecycleCtx.Done():
153+
case<-s.lifecycleCtx.Done():
154154
returntrue
155155
default:
156156
returnfalse
157157
}
158158
}
159159

160160
// Shutdown waits for all exiting in-flight requests to complete, or the context to expire, whichever comes first.
161-
func (d*Server)Shutdown(ctx context.Context)error {
161+
func (s*Server)Shutdown(ctx context.Context)error {
162162
varerrerror
163-
d.shutdownOnce.Do(func() {
164-
d.cancelFn()
163+
s.shutdownOnce.Do(func() {
164+
s.cancelFn()
165165

166166
// Wait for any outstanding connections to terminate.
167-
d.wg.Wait()
167+
s.wg.Wait()
168168

169169
select {
170170
case<-ctx.Done():
171-
d.logger.Warn(ctx,"graceful shutdown failed",slog.Error(ctx.Err()))
171+
s.logger.Warn(ctx,"graceful shutdown failed",slog.Error(ctx.Err()))
172172
err=ctx.Err()
173173
return
174174
default:
175175
}
176176

177-
d.logger.Info(ctx,"shutting down request pool")
178-
iferr=d.requestBridgePool.Shutdown(ctx);err!=nil {
179-
d.logger.Error(ctx,"request pool shutdown failed with error",slog.Error(err))
177+
s.logger.Info(ctx,"shutting down request pool")
178+
iferr=s.requestBridgePool.Shutdown(ctx);err!=nil {
179+
s.logger.Error(ctx,"request pool shutdown failed with error",slog.Error(err))
180180
}
181181

182-
d.logger.Info(ctx,"gracefully shutdown")
182+
s.logger.Info(ctx,"gracefully shutdown")
183183
})
184184
returnerr
185185
}

‎enterprise/x/aibridged/aibridged_test.go‎

Lines changed: 18 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,10 @@ func newTestServer(t *testing.T) (*aibridged.Server, *mock.MockDRPCClient, *mock
3232
client:=mock.NewMockDRPCClient(ctrl)
3333
pool:=mock.NewMockPooler(ctrl)
3434

35+
conn:=&mockDRPCConn{}
36+
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
37+
pool.EXPECT().Shutdown(gomock.Any()).MinTimes(1).Return(nil)
38+
3539
srv,err:=aibridged.New(
3640
t.Context(),
3741
pool,
@@ -40,6 +44,9 @@ func newTestServer(t *testing.T) (*aibridged.Server, *mock.MockDRPCClient, *mock
4044
},
4145
logger)
4246
require.NoError(t,err,"create new aibridged")
47+
t.Cleanup(func() {
48+
srv.Shutdown(context.Background())
49+
})
4350

4451
returnsrv,client,pool
4552
}
@@ -115,7 +122,7 @@ func TestServeHTTP_FailureModes(t *testing.T) {
115122
// Should pass authorization.
116123
client.EXPECT().IsAuthorized(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId:uuid.NewString()},nil)
117124
// But fail when acquiring a pool instance.
118-
pool.EXPECT().Acquire(gomock.Any(),gomock.Any(),gomock.Any()).AnyTimes().Return(nil,xerrors.New("oops"))
125+
pool.EXPECT().Acquire(gomock.Any(),gomock.Any(),gomock.Any(),gomock.Any()).AnyTimes().Return(nil,xerrors.New("oops"))
119126
},
120127
expectedErr:aibridged.ErrAcquireRequestHandler,
121128
expectedStatus:http.StatusInternalServerError,
@@ -228,38 +235,6 @@ func (*mockHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
228235
_,_=rw.Write([]byte(r.URL.Path))
229236
}
230237

231-
// TestPoolHandler validates that an http.Handler can be acquired from a given [aibridged.Pooler]
232-
// and have its HTTP handler invoked.
233-
//
234-
// We're not actually testing the routing, since that is being tested by [aibridge.RequestBridge].
235-
//
236-
// We're validating that a request can be successfully processed by aibridged
237-
// (i.e. authn/z, acquire pool instance) and what happens thereafter is a black box to aibridged.
238-
funcTestPoolHandler(t*testing.T) {
239-
t.Parallel()
240-
241-
srv,client,pool:=newTestServer(t)
242-
243-
conn:=&mockDRPCConn{}
244-
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
245-
// Authorize all requests.
246-
client.EXPECT().IsAuthorized(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId:uuid.NewString()},nil)
247-
pool.EXPECT().Acquire(gomock.Any(),gomock.Any(),gomock.Any()).AnyTimes().Return(&mockHandler{},nil)
248-
249-
ctx:=testutil.Context(t,testutil.WaitShort)
250-
path:="/irrelevant"
251-
req,err:=http.NewRequestWithContext(ctx,http.MethodPost,path,nil)
252-
require.NoError(t,err,"make request to test server")
253-
req.Header.Add("Authorization","Bearer key")
254-
255-
rec:=httptest.NewRecorder()
256-
srv.ServeHTTP(rec,req)
257-
258-
require.Equal(t,http.StatusOK,rec.Code)
259-
require.NotNil(t,rec.Body)
260-
require.Equal(t,path,rec.Body.String())
261-
}
262-
263238
// TestRouting validates that a request which originates with aibridged will be handled
264239
// by coder/aibridge's handling logic in a provider-specific manner.
265240
// We must validate that logic that pertains to coder/coder is exercised.
@@ -284,13 +259,13 @@ func TestRouting(t *testing.T) {
284259
{
285260
name:"openai chat completions",
286261
path:"/openai/v1/chat/completions",
287-
expectedStatus:http.StatusOK,
262+
expectedStatus:http.StatusTeapot,// Nonsense status to indicate server was hit.
288263
expectedHits:1,
289264
},
290265
{
291266
name:"anthropic messages",
292267
path:"/anthropic/v1/messages",
293-
expectedStatus:http.StatusOK,
268+
expectedStatus:http.StatusTeapot,// Nonsense status to indicate server was hit.
294269
expectedHits:1,
295270
},
296271
}
@@ -310,14 +285,12 @@ func TestRouting(t *testing.T) {
310285
logger:=slogtest.Make(t,&slogtest.Options{IgnoreErrors:true})
311286
ctrl:=gomock.NewController(t)
312287
client:=mock.NewMockDRPCClient(ctrl)
313-
pool,err:=aibridged.NewCachedBridgePool(10, aibridge.Config{
314-
OpenAI: aibridge.ProviderConfig{
315-
BaseURL:openaiSrv.URL,
316-
},
317-
Anthropic: aibridge.ProviderConfig{
318-
BaseURL:antSrv.URL,
319-
},
320-
},logger)
288+
289+
providers:= []aibridge.Provider{
290+
aibridge.NewOpenAIProvider(aibridge.ProviderConfig{BaseURL:openaiSrv.URL}),
291+
aibridge.NewAnthropicProvider(aibridge.ProviderConfig{BaseURL:antSrv.URL}),
292+
}
293+
pool,err:=aibridged.NewCachedBridgePool(aibridged.DefaultPoolOptions,providers,logger)
321294
require.NoError(t,err)
322295
conn:=&mockDRPCConn{}
323296
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
@@ -331,9 +304,6 @@ func TestRouting(t *testing.T) {
331304
interceptionID=in.GetId()
332305
return&proto.RecordInterceptionResponse{},nil
333306
})
334-
client.EXPECT().RecordPromptUsage(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.RecordPromptUsageResponse{},nil)
335-
client.EXPECT().RecordTokenUsage(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.RecordTokenUsageResponse{},nil)
336-
client.EXPECT().RecordToolUsage(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.RecordToolUsageResponse{},nil)
337307

338308
// Given: aibridged is started.
339309
srv,err:=aibridged.New(t.Context(),pool,func(ctx context.Context) (aibridged.DRPCClient,error) {
@@ -356,6 +326,8 @@ func TestRouting(t *testing.T) {
356326
srv.ServeHTTP(rec,req)
357327

358328
// Then: the upstream server will have received a number of hits.
329+
// NOTE: we *expect* the interceptions to fail because [mockAIUpstreamServer] returns a nonsense status code.
330+
// We only need to test that the request was routed, NOT processed.
359331
require.Equal(t,tc.expectedStatus,rec.Code)
360332
assert.EqualValues(t,tc.expectedHits,upstreamSrv.Hits())
361333
iftc.expectedHits>0 {

‎enterprise/x/aibridged/aibridgedmock/poolmock.go‎

Lines changed: 4 additions & 4 deletions
Some generated files are not rendered by default. Learn more aboutcustomizing how changed files appear on GitHub.

‎enterprise/x/aibridged/http.go‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -30,10 +30,10 @@ var (
3030
//
3131
// A [DRPCClient] is provided to the [aibridge.RequestBridge] instance so that data can
3232
// be passed up to a [DRPCServer] for persistence.
33-
func (d*Server)ServeHTTP(rw http.ResponseWriter,r*http.Request) {
33+
func (s*Server)ServeHTTP(rw http.ResponseWriter,r*http.Request) {
3434
ctx:=r.Context()
3535

36-
logger:=d.logger.With(slog.F("path",r.URL.Path))
36+
logger:=s.logger.With(slog.F("path",r.URL.Path))
3737

3838
key:=strings.TrimSpace(ExtractAuthToken(r.Header))
3939
ifkey=="" {
@@ -42,7 +42,7 @@ func (d *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
4242
return
4343
}
4444

45-
client,err:=d.Client()
45+
client,err:=s.Client()
4646
iferr!=nil {
4747
logger.Warn(ctx,"failed to connect to coderd",slog.Error(err))
4848
http.Error(rw,ErrConnect.Error(),http.StatusServiceUnavailable)
@@ -66,7 +66,7 @@ func (d *Server) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
6666
return
6767
}
6868

69-
handler,err:=d.GetRequestHandler(ctx,Request{
69+
handler,err:=s.GetRequestHandler(ctx,Request{
7070
SessionKey:key,
7171
InitiatorID:id,
7272
})

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp