@@ -15,6 +15,7 @@ import (
1515 "strings"
1616 "testing"
1717
18+ "github.com/Wei-Shaw/sub2api/internal/config"
1819 infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
1920 "github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
2021 "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
@@ -29,6 +30,20 @@ const (
2930 userAgentIdentityPublicGroupID int64 = 8202
3031)
3132
33+ type userAccountTestModelCatalog struct {}
34+
35+ func (userAccountTestModelCatalog ) ListPricedModelIDs (context.Context , []string ) ([]string , error ) {
36+ return nil , nil
37+ }
38+
39+ func (userAccountTestModelCatalog ) ListSelectablePricedModelIDs (context.Context , service.PricedModelQuery ) ([]string , error ) {
40+ return []string {"selected-model" , "gpt-new" }, nil
41+ }
42+
43+ func (userAccountTestModelCatalog ) IsModelPriced (context.Context , service.PricedModelQuery , string ) (bool , error ) {
44+ return false , nil
45+ }
46+
3247type userAgentIdentityShareRepo struct {
3348 service.AccountRepository
3449 accounts map [int64 ]* service.Account
@@ -259,6 +274,7 @@ type userAgentIdentityValidationUpstream struct {
259274 body string
260275 calls int
261276 lastAuthorization string
277+ lastModel string
262278}
263279
264280func (u * userAgentIdentityValidationUpstream ) Do (req * http.Request , _ string , _ int64 , _ int ) (* http.Response , error ) {
@@ -272,6 +288,13 @@ func (u *userAgentIdentityValidationUpstream) DoWithTLS(req *http.Request, _ str
272288func (u * userAgentIdentityValidationUpstream ) response (req * http.Request ) * http.Response {
273289 u .calls ++
274290 u .lastAuthorization = req .Header .Get ("Authorization" )
291+ if req .Body != nil {
292+ var payload struct {
293+ Model string `json:"model"`
294+ }
295+ _ = json .NewDecoder (req .Body ).Decode (& payload )
296+ u .lastModel = payload .Model
297+ }
275298 statusCode := u .statusCode
276299 if statusCode == 0 {
277300 statusCode = http .StatusOK
@@ -368,7 +391,8 @@ func newUserAgentIdentityShareHandler(
368391 accountService .SetAccountShareModeRepository (placementRepo )
369392 accountService .SetAgentIdentityWSInvalidator (invalidatorProxy )
370393 upstream := & userAgentIdentityValidationUpstream {statusCode : upstreamStatus , body : upstreamBody }
371- accountTestService := service .NewAccountTestService (repo , nil , nil , nil , upstream , nil , nil , nil , invalidatorProxy )
394+ accountTestService := service .NewAccountTestService (repo , nil , nil , nil , upstream , & config.Config {}, nil , nil , invalidatorProxy )
395+ accountTestService .SetModelResolver (service .NewAccountTestModelResolver (userAccountTestModelCatalog {}))
372396 handler := NewUserAccountHandler (accountService , nil , accountTestService , nil , nil , nil , nil , nil , nil , nil )
373397 return handler , repo , upstream , invalidator , placementRepo
374398}
@@ -389,10 +413,10 @@ func runUserAgentIdentityUpdateRequest(t *testing.T, handler *UserAccountHandler
389413 return recorder
390414}
391415
392- func TestUserAccountHandlerTestRejectsModelOutsideOwnerWhitelist (t * testing.T ) {
416+ func TestUserAccountHandlerTestRejectsModelOutsideOwnerWhitelistForPublicAccount (t * testing.T ) {
393417 gin .SetMode (gin .TestMode )
394418 ownerUserID := int64 (101 )
395- account := newUserAgentIdentityShareAccount (t , ownerUserID , service .AccountShareModePrivate , service .AccountShareStatusApproved )
419+ account := newUserAgentIdentityShareAccount (t , ownerUserID , service .AccountShareModePublic , service .AccountShareStatusApproved )
396420 account .Credentials ["model_mapping" ] = map [string ]any {"selected-model" : "selected-model" }
397421 handler , _ , upstream , _ , _ := newUserAgentIdentityShareHandler (t , account , http .StatusOK , "" )
398422
@@ -412,6 +436,53 @@ func TestUserAccountHandlerTestRejectsModelOutsideOwnerWhitelist(t *testing.T) {
412436 require .Zero (t , upstream .calls )
413437}
414438
439+ func TestUserAccountHandlerTestPrivateInheritanceAllowsPricedModelOutsideLegacyMapping (t * testing.T ) {
440+ gin .SetMode (gin .TestMode )
441+ ownerUserID := int64 (101 )
442+ account := newUserAgentIdentityShareAccount (t , ownerUserID , service .AccountShareModePrivate , service .AccountShareStatusApproved )
443+ account .Type = service .AccountTypeAPIKey
444+ account .Credentials = map [string ]any {"api_key" : "test-key" , "model_mapping" : map [string ]any {"old-model" : "old-model" }}
445+ handler , _ , upstream , _ , _ := newUserAgentIdentityShareHandler (t , account , http .StatusOK , "" )
446+
447+ router := gin .New ()
448+ router .POST ("/accounts/:id/test" , func (c * gin.Context ) {
449+ c .Set (string (middleware2 .ContextKeyUser ), middleware2.AuthSubject {UserID : ownerUserID })
450+ handler .Test (c )
451+ })
452+ recorder := httptest .NewRecorder ()
453+ request := httptest .NewRequest (http .MethodPost , "/accounts/1/test" , strings .NewReader (`{"model_id":"gpt-new"}` ))
454+ request .Header .Set ("Content-Type" , "application/json" )
455+ router .ServeHTTP (recorder , request )
456+
457+ require .Equal (t , http .StatusOK , recorder .Code , recorder .Body .String ())
458+ require .Contains (t , recorder .Body .String (), `"success":true` )
459+ require .Equal (t , 1 , upstream .calls )
460+ require .Equal (t , "gpt-new" , upstream .lastModel )
461+ }
462+
463+ func TestUserAccountHandlerTestPrivateInheritanceRejectsUnpricedModel (t * testing.T ) {
464+ gin .SetMode (gin .TestMode )
465+ ownerUserID := int64 (101 )
466+ account := newUserAgentIdentityShareAccount (t , ownerUserID , service .AccountShareModePrivate , service .AccountShareStatusApproved )
467+ account .Type = service .AccountTypeAPIKey
468+ account .Credentials = map [string ]any {"api_key" : "test-key" }
469+ handler , _ , upstream , _ , _ := newUserAgentIdentityShareHandler (t , account , http .StatusOK , "" )
470+
471+ router := gin .New ()
472+ router .POST ("/accounts/:id/test" , func (c * gin.Context ) {
473+ c .Set (string (middleware2 .ContextKeyUser ), middleware2.AuthSubject {UserID : ownerUserID })
474+ handler .Test (c )
475+ })
476+ recorder := httptest .NewRecorder ()
477+ request := httptest .NewRequest (http .MethodPost , "/accounts/1/test" , strings .NewReader (`{"model_id":"unpriced-model"}` ))
478+ request .Header .Set ("Content-Type" , "application/json" )
479+ router .ServeHTTP (recorder , request )
480+
481+ require .Equal (t , http .StatusBadRequest , recorder .Code , recorder .Body .String ())
482+ require .Contains (t , recorder .Body .String (), "ACCOUNT_TEST_MODEL_NOT_AVAILABLE" )
483+ require .Zero (t , upstream .calls )
484+ }
485+
415486func TestIsOpenAIUsageLimitReachedValidationError (t * testing.T ) {
416487 require .True (t , isOpenAIUsageLimitReachedValidationError (`API returned 429: {"error":{"type":"usage_limit_reached","message":"The usage limit has been reached"}}` ))
417488 require .True (t , isOpenAIUsageLimitReachedValidationError (`API returned 429: {"error": {"type": "usage_limit_reached"}}` ))
0 commit comments