Commit 7dc7ff22 authored by IanShaw027's avatar IanShaw027
Browse files

fix: preserve messages dispatch config in repository hydration

parent 67a05dfc
...@@ -167,6 +167,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se ...@@ -167,6 +167,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
group.FieldSupportedModelScopes, group.FieldSupportedModelScopes,
group.FieldAllowMessagesDispatch, group.FieldAllowMessagesDispatch,
group.FieldDefaultMappedModel, group.FieldDefaultMappedModel,
group.FieldMessagesDispatchModelConfig,
) )
}). }).
Only(ctx) Only(ctx)
...@@ -689,6 +690,7 @@ func groupEntityToService(g *dbent.Group) *service.Group { ...@@ -689,6 +690,7 @@ func groupEntityToService(g *dbent.Group) *service.Group {
RequireOAuthOnly: g.RequireOauthOnly, RequireOAuthOnly: g.RequireOauthOnly,
RequirePrivacySet: g.RequirePrivacySet, RequirePrivacySet: g.RequirePrivacySet,
DefaultMappedModel: g.DefaultMappedModel, DefaultMappedModel: g.DefaultMappedModel,
MessagesDispatchModelConfig: g.MessagesDispatchModelConfig,
CreatedAt: g.CreatedAt, CreatedAt: g.CreatedAt,
UpdatedAt: g.UpdatedAt, UpdatedAt: g.UpdatedAt,
} }
......
...@@ -86,6 +86,45 @@ func (s *APIKeyRepoSuite) TestGetByKey_NotFound() { ...@@ -86,6 +86,45 @@ func (s *APIKeyRepoSuite) TestGetByKey_NotFound() {
s.Require().Error(err, "expected error for non-existent key") s.Require().Error(err, "expected error for non-existent key")
} }
func (s *APIKeyRepoSuite) TestGetByKeyForAuth_PreservesMessagesDispatchModelConfig() {
user := s.mustCreateUser("getbykey-auth-dispatch@test.com")
group, err := s.client.Group.Create().
SetName("g-auth-dispatch").
SetPlatform(service.PlatformOpenAI).
SetStatus(service.StatusActive).
SetSubscriptionType(service.SubscriptionTypeStandard).
SetRateMultiplier(1).
SetAllowMessagesDispatch(true).
SetDefaultMappedModel("gpt-5.4").
SetMessagesDispatchModelConfig(service.OpenAIMessagesDispatchModelConfig{
OpusMappedModel: "gpt-5.4-nano",
SonnetMappedModel: "gpt-5.3-codex",
HaikuMappedModel: "gpt-5.4-mini",
ExactModelMappings: map[string]string{
"claude-sonnet-4.5": "gpt-5.4-nano",
},
}).
Save(s.ctx)
s.Require().NoError(err)
key := &service.APIKey{
UserID: user.ID,
Key: "sk-getbykey-auth-dispatch",
Name: "Dispatch Key",
GroupID: &group.ID,
Status: service.StatusActive,
}
s.Require().NoError(s.repo.Create(s.ctx, key))
got, err := s.repo.GetByKeyForAuth(s.ctx, key.Key)
s.Require().NoError(err)
s.Require().NotNil(got.Group)
s.Require().True(got.Group.AllowMessagesDispatch)
s.Require().Equal("gpt-5.4", got.Group.DefaultMappedModel)
s.Require().Equal("gpt-5.4-nano", got.Group.MessagesDispatchModelConfig.OpusMappedModel)
s.Require().Equal("gpt-5.4-nano", got.Group.MessagesDispatchModelConfig.ExactModelMappings["claude-sonnet-4.5"])
}
// --- Update --- // --- Update ---
func (s *APIKeyRepoSuite) TestUpdate() { func (s *APIKeyRepoSuite) TestUpdate() {
......
package repository
import (
"context"
"testing"
dbent "github.com/Wei-Shaw/sub2api/ent"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/stretchr/testify/require"
)
func TestGroupEntityToService_PreservesMessagesDispatchModelConfig(t *testing.T) {
group := &dbent.Group{
ID: 1,
Name: "openai-dispatch",
Platform: service.PlatformOpenAI,
Status: service.StatusActive,
SubscriptionType: service.SubscriptionTypeStandard,
RateMultiplier: 1,
AllowMessagesDispatch: true,
DefaultMappedModel: "gpt-5.4",
MessagesDispatchModelConfig: service.OpenAIMessagesDispatchModelConfig{
OpusMappedModel: "gpt-5.4-nano",
SonnetMappedModel: "gpt-5.3-codex",
HaikuMappedModel: "gpt-5.4-mini",
ExactModelMappings: map[string]string{
"claude-sonnet-4.5": "gpt-5.4-nano",
},
},
}
got := groupEntityToService(group)
require.NotNil(t, got)
require.Equal(t, group.MessagesDispatchModelConfig, got.MessagesDispatchModelConfig)
}
func TestAPIKeyRepository_GetByKeyForAuth_PreservesMessagesDispatchModelConfig_SQLite(t *testing.T) {
repo, client := newAPIKeyRepoSQLite(t)
ctx := context.Background()
user := mustCreateAPIKeyRepoUser(t, ctx, client, "getbykey-auth-dispatch-unit@test.com")
group, err := client.Group.Create().
SetName("g-auth-dispatch-unit").
SetPlatform(service.PlatformOpenAI).
SetStatus(service.StatusActive).
SetSubscriptionType(service.SubscriptionTypeStandard).
SetRateMultiplier(1).
SetAllowMessagesDispatch(true).
SetDefaultMappedModel("gpt-5.4").
SetMessagesDispatchModelConfig(service.OpenAIMessagesDispatchModelConfig{
OpusMappedModel: "gpt-5.4-nano",
SonnetMappedModel: "gpt-5.3-codex",
HaikuMappedModel: "gpt-5.4-mini",
ExactModelMappings: map[string]string{
"claude-sonnet-4.5": "gpt-5.4-nano",
},
}).
Save(ctx)
require.NoError(t, err)
key := &service.APIKey{
UserID: user.ID,
Key: "sk-getbykey-auth-dispatch-unit",
Name: "Dispatch Key Unit",
GroupID: &group.ID,
Status: service.StatusActive,
}
require.NoError(t, repo.Create(ctx, key))
got, err := repo.GetByKeyForAuth(ctx, key.Key)
require.NoError(t, err)
require.NotNil(t, got.Group)
require.Equal(t, group.MessagesDispatchModelConfig, got.Group.MessagesDispatchModelConfig)
}
...@@ -113,6 +113,33 @@ func (s *GroupRepoSuite) TestUpdate() { ...@@ -113,6 +113,33 @@ func (s *GroupRepoSuite) TestUpdate() {
s.Require().Equal("updated", got.Name) s.Require().Equal("updated", got.Name)
} }
func (s *GroupRepoSuite) TestGetByID_PreservesMessagesDispatchModelConfig() {
group := &service.Group{
Name: "openai-dispatch",
Platform: service.PlatformOpenAI,
RateMultiplier: 1.0,
IsExclusive: false,
Status: service.StatusActive,
SubscriptionType: service.SubscriptionTypeStandard,
AllowMessagesDispatch: true,
DefaultMappedModel: "gpt-5.4",
MessagesDispatchModelConfig: service.OpenAIMessagesDispatchModelConfig{
OpusMappedModel: "gpt-5.4",
SonnetMappedModel: "gpt-5.3-codex",
HaikuMappedModel: "gpt-5.4-mini",
ExactModelMappings: map[string]string{
"claude-sonnet-4.5": "gpt-5.4-nano",
},
},
}
s.Require().NoError(s.repo.Create(s.ctx, group))
got, err := s.repo.GetByID(s.ctx, group.ID)
s.Require().NoError(err)
s.Require().Equal(group.MessagesDispatchModelConfig, got.MessagesDispatchModelConfig)
}
func (s *GroupRepoSuite) TestDelete() { func (s *GroupRepoSuite) TestDelete() {
group := &service.Group{ group := &service.Group{
Name: "to-delete", Name: "to-delete",
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment