package routes import ( "bytes" "context" "io" "mime/multipart" "net/http" "net/http/httptest" "strings" "testing" servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware" "github.com/Wei-Shaw/sub2api/internal/service" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) type compositeRouteRepoStub struct { routes []service.CompositeModelRoute } func (s compositeRouteRepoStub) ListByGroup(ctx context.Context, groupID int64, includeDisabled bool) ([]service.CompositeModelRoute, error) { routes := make([]service.CompositeModelRoute, 0, len(s.routes)) for _, route := range s.routes { if route.GroupID != groupID { continue } if !includeDisabled && !route.Enabled { continue } routes = append(routes, route) } return routes, nil } func (s compositeRouteRepoStub) Create(ctx context.Context, route *service.CompositeModelRoute) error { return nil } func (s compositeRouteRepoStub) Update(ctx context.Context, route *service.CompositeModelRoute) error { return nil } func (s compositeRouteRepoStub) Delete(ctx context.Context, id int64) error { return nil } func (s compositeRouteRepoStub) DeleteByGroup(ctx context.Context, groupID int64) error { return nil } func TestCompositeTargetPlatformMiddlewareResolvesModelAndRestoresBody(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) { groupID := int64(1) c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{ GroupID: &groupID, Group: &service.Group{Platform: service.PlatformComposite}, }) c.Next() }))) router.Use(compositeTargetPlatformMiddleware(nil)) router.POST("/", func(c *gin.Context) { platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, service.PlatformOpenAI, platform) body, err := io.ReadAll(c.Request.Body) require.NoError(t, err) require.JSONEq(t, `{"model":"gpt-5"}`, string(body)) c.Status(http.StatusNoContent) }) req := httptest.NewRequest(http.MethodPost, "/", strings.NewReader(`{"model":"gpt-5"}`)) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusNoContent, w.Code) } func TestCompositeTargetPlatformMiddlewareUsesExplicitRouteAndRewritesBody(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{ routes: []service.CompositeModelRoute{ { ID: 1, GroupID: 1, PublicModel: "openrouter/gpt-5", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, UpstreamModel: "gpt-5", Endpoint: service.CompositeRouteEndpointAny, Priority: 100, Enabled: true, }, }, }) router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) { groupID := int64(1) c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{ GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, }) c.Next() }))) router.Use(compositeTargetPlatformMiddleware(resolver)) router.POST("/v1/chat/completions", func(c *gin.Context) { platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, service.PlatformOpenAI, platform) upstreamModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, "gpt-5", upstreamModel) body, err := io.ReadAll(c.Request.Body) require.NoError(t, err) require.JSONEq(t, `{"model":"gpt-5","messages":[]}`, string(body)) c.Status(http.StatusNoContent) }) req := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"openrouter/gpt-5","messages":[]}`)) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusNoContent, w.Code) } func TestCompositeTargetPlatformMiddlewareRewritesNestedLiveModel(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{ routes: []service.CompositeModelRoute{ { ID: 1, GroupID: 1, PublicModel: "live-alias", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, UpstreamModel: "gpt-live", Endpoint: service.CompositeRouteEndpointAny, Priority: 100, Enabled: true, }, }, }) router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) { groupID := int64(1) c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{ GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, }) c.Next() }))) router.Use(compositeTargetPlatformMiddleware(resolver)) router.POST("/backend-api/codex/realtime/calls", func(c *gin.Context) { platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, service.PlatformOpenAI, platform) body, err := io.ReadAll(c.Request.Body) require.NoError(t, err) require.JSONEq(t, `{"session":{"model":"gpt-live"},"sdp":"v=0"}`, string(body)) c.Status(http.StatusNoContent) }) req := httptest.NewRequest( http.MethodPost, "/backend-api/codex/realtime/calls", strings.NewReader(`{"session":{"model":"live-alias"},"sdp":"v=0"}`), ) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusNoContent, w.Code) } func TestCompositeRequestModelFromMultipartLiveSession(t *testing.T) { var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("sdp", "v=0")) require.NoError(t, writer.WriteField("session", `{"model":"live-alias"}`)) require.NoError(t, writer.Close()) require.Equal(t, "live-alias", compositeRequestModelFromBody(writer.FormDataContentType(), body.Bytes())) } func TestCompositeCodexControlPathsUseResponsesRoutes(t *testing.T) { for _, path := range []string{ "/v1/alpha/search", "/backend-api/codex/alpha/search", "/v1/live", "/backend-api/codex/realtime/calls", } { require.Equal(t, service.CompositeRouteEndpointResponses, compositeRouteEndpointForPath(path), "path=%s", path) } } func TestCompositeTargetPlatformMiddlewareUsesExplicitRouteForMultipartImages(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{ routes: []service.CompositeModelRoute{ { ID: 1, GroupID: 1, PublicModel: "image-alias", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformOpenAI, UpstreamModel: "gpt-image-1", Endpoint: service.CompositeRouteEndpointImages, Priority: 100, Enabled: true, }, }, }) router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) { groupID := int64(1) c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{ GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, }) c.Next() }))) router.Use(compositeTargetPlatformMiddleware(resolver)) router.POST("/v1/images/edits", func(c *gin.Context) { platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, service.PlatformOpenAI, platform) upstreamModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, "gpt-image-1", upstreamModel) publicModel, ok := service.RequestedPublicModelFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, "image-alias", publicModel) body, err := io.ReadAll(c.Request.Body) require.NoError(t, err) require.Contains(t, string(body), "image-alias") c.Status(http.StatusNoContent) }) var body bytes.Buffer writer := multipart.NewWriter(&body) require.NoError(t, writer.WriteField("model", "image-alias")) require.NoError(t, writer.WriteField("prompt", "draw")) require.NoError(t, writer.Close()) req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(body.Bytes())) req.Header.Set("Content-Type", writer.FormDataContentType()) w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusNoContent, w.Code) } func TestCompositeGeminiTargetPlatformMiddlewareUsesPathRoute(t *testing.T) { gin.SetMode(gin.TestMode) router := gin.New() resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{ routes: []service.CompositeModelRoute{ { ID: 1, GroupID: 1, PublicModel: "openrouter/gemini-pro", MatchType: service.CompositeRouteMatchExact, TargetPlatform: service.PlatformGemini, UpstreamModel: "gemini-2.5-pro", Endpoint: service.CompositeRouteEndpointGemini, Priority: 100, Enabled: true, }, }, }) router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) { groupID := int64(1) c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{ GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformComposite}, }) c.Next() }))) router.Use(compositeGeminiTargetPlatformMiddleware(resolver)) router.POST("/v1beta/models/*modelAction", func(c *gin.Context) { platform, ok := service.ResolvedTargetPlatformFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, service.PlatformGemini, platform) upstreamModel, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context()) require.True(t, ok) require.Equal(t, "gemini-2.5-pro", upstreamModel) c.Status(http.StatusNoContent) }) req := httptest.NewRequest(http.MethodPost, "/v1beta/models/openrouter/gemini-pro:generateContent", strings.NewReader(`{"contents":[]}`)) req.Header.Set("Content-Type", "application/json") w := httptest.NewRecorder() router.ServeHTTP(w, req) require.Equal(t, http.StatusNoContent, w.Code) }