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

Commit934108a

Browse files
dannykoppingdeansheather
authored andcommitted
chore: more tests
1 parentf69a5d2 commit934108a

File tree

9 files changed

+422
-269
lines changed

9 files changed

+422
-269
lines changed

‎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
}

‎aibridged/aibridged_test.go‎

Lines changed: 20 additions & 49 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
}
@@ -53,6 +60,7 @@ func (*mockDRPCConn) Transport() drpc.Transport { return nil }
5360
func (*mockDRPCConn)Invoke(ctx context.Context,rpcstring,enc drpc.Encoding,in,out drpc.Message)error {
5461
returnnil
5562
}
63+
5664
func (*mockDRPCConn)NewStream(ctx context.Context,rpcstring,enc drpc.Encoding) (drpc.Stream,error) {
5765
// nolint:nilnil // Chillchill.
5866
returnnil,nil
@@ -61,9 +69,7 @@ func (*mockDRPCConn) NewStream(ctx context.Context, rpc string, enc drpc.Encodin
6169
funcTestServeHTTP_FailureModes(t*testing.T) {
6270
t.Parallel()
6371

64-
var (
65-
defaultHeaders=map[string]string{"Authorization":"Bearer key"}
66-
)
72+
defaultHeaders:=map[string]string{"Authorization":"Bearer key"}
6773

6874
cases:= []struct {
6975
namestring
@@ -116,7 +122,7 @@ func TestServeHTTP_FailureModes(t *testing.T) {
116122
// Should pass authorization.
117123
client.EXPECT().IsAuthorized(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId:uuid.NewString()},nil)
118124
// But fail when acquiring a pool instance.
119-
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"))
120126
},
121127
expectedErr:aibridged.ErrAcquireRequestHandler,
122128
expectedStatus:http.StatusInternalServerError,
@@ -229,38 +235,6 @@ func (*mockHandler) ServeHTTP(rw http.ResponseWriter, r *http.Request) {
229235
_,_=rw.Write([]byte(r.URL.Path))
230236
}
231237

232-
// TestPoolHandler validates that an http.Handler can be acquired from a given [aibridged.Pooler]
233-
// and have its HTTP handler invoked.
234-
//
235-
// We're not actually testing the routing, since that is being tested by [aibridge.RequestBridge].
236-
//
237-
// We're validating that a request can be successfully processed by aibridged
238-
// (i.e. authn/z, acquire pool instance) and what happens thereafter is a black box to aibridged.
239-
funcTestPoolHandler(t*testing.T) {
240-
t.Parallel()
241-
242-
srv,client,pool:=newTestServer(t)
243-
244-
conn:=&mockDRPCConn{}
245-
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
246-
// Authorize all requests.
247-
client.EXPECT().IsAuthorized(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.IsAuthorizedResponse{OwnerId:uuid.NewString()},nil)
248-
pool.EXPECT().Acquire(gomock.Any(),gomock.Any(),gomock.Any()).AnyTimes().Return(&mockHandler{},nil)
249-
250-
ctx:=testutil.Context(t,testutil.WaitShort)
251-
path:="/irrelevant"
252-
req,err:=http.NewRequestWithContext(ctx,http.MethodPost,path,nil)
253-
require.NoError(t,err,"make request to test server")
254-
req.Header.Add("Authorization","Bearer key")
255-
256-
rec:=httptest.NewRecorder()
257-
srv.ServeHTTP(rec,req)
258-
259-
require.Equal(t,http.StatusOK,rec.Code)
260-
require.NotNil(t,rec.Body)
261-
require.Equal(t,path,rec.Body.String())
262-
}
263-
264238
// TestRouting validates that a request which originates with aibridged will be handled
265239
// by coder/aibridge's handling logic in a provider-specific manner.
266240
// We must validate that logic that pertains to coder/coder is exercised.
@@ -285,13 +259,13 @@ func TestRouting(t *testing.T) {
285259
{
286260
name:"openai chat completions",
287261
path:"/openai/v1/chat/completions",
288-
expectedStatus:http.StatusOK,
262+
expectedStatus:http.StatusTeapot,// Nonsense status to indicate server was hit.
289263
expectedHits:1,
290264
},
291265
{
292266
name:"anthropic messages",
293267
path:"/anthropic/v1/messages",
294-
expectedStatus:http.StatusOK,
268+
expectedStatus:http.StatusTeapot,// Nonsense status to indicate server was hit.
295269
expectedHits:1,
296270
},
297271
}
@@ -311,14 +285,12 @@ func TestRouting(t *testing.T) {
311285
logger:=slogtest.Make(t,&slogtest.Options{IgnoreErrors:true})
312286
ctrl:=gomock.NewController(t)
313287
client:=mock.NewMockDRPCClient(ctrl)
314-
pool,err:=aibridged.NewCachedBridgePool(10, aibridge.Config{
315-
OpenAI: aibridge.ProviderConfig{
316-
BaseURL:openaiSrv.URL,
317-
},
318-
Anthropic: aibridge.ProviderConfig{
319-
BaseURL:antSrv.URL,
320-
},
321-
},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)
322294
require.NoError(t,err)
323295
conn:=&mockDRPCConn{}
324296
client.EXPECT().DRPCConn().AnyTimes().Return(conn)
@@ -332,9 +304,6 @@ func TestRouting(t *testing.T) {
332304
interceptionID=in.GetId()
333305
return&proto.RecordInterceptionResponse{},nil
334306
})
335-
client.EXPECT().RecordPromptUsage(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.RecordPromptUsageResponse{},nil)
336-
client.EXPECT().RecordTokenUsage(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.RecordTokenUsageResponse{},nil)
337-
client.EXPECT().RecordToolUsage(gomock.Any(),gomock.Any()).AnyTimes().Return(&proto.RecordToolUsageResponse{},nil)
338307

339308
// Given: aibridged is started.
340309
srv,err:=aibridged.New(t.Context(),pool,func(ctx context.Context) (aibridged.DRPCClient,error) {
@@ -357,6 +326,8 @@ func TestRouting(t *testing.T) {
357326
srv.ServeHTTP(rec,req)
358327

359328
// 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.
360331
require.Equal(t,tc.expectedStatus,rec.Code)
361332
assert.EqualValues(t,tc.expectedHits,upstreamSrv.Hits())
362333
iftc.expectedHits>0 {

‎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.

‎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