new-api/relay/helper/dimension_resolver_test.go

package helper

import (
	"bytes"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/QuantumNous/new-api/types"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
)

func TestDimensionResolver_JSONBody(t *testing.T) {
	c := newPricingJSONContext(`{"resolution":"720P","duration":"10","options":{"size":"1280x720"},"image_url":"https://example.test/a.png"}`)

	resolver, err := NewDimensionResolver(c)
	require.NoError(t, err)

	values, err := resolver.ResolveDimensions([]types.PricingDimension{
		{Key: "resolution", Source: "request.resolution", Type: "string"},
		{Key: "size", Source: "request.options.size", Type: "string"},
		{Key: "duration", Source: "request.duration", Type: "number"},
		{Key: "mode", Source: "derived.input_mode", Type: "string"},
	})
	require.NoError(t, err)

	require.Equal(t, "720P", values.Values["resolution"])
	require.Equal(t, "1280x720", values.Values["size"])
	require.Equal(t, 10.0, values.Values["duration"])
	require.Equal(t, "i2v", values.Values["mode"])
	require.Empty(t, values.Missing)
}

func TestDimensionResolver_DerivedInputModePriority(t *testing.T) {
	c := newPricingJSONContext(`{"image_url":"https://example.test/a.png","video_url":"https://example.test/a.mp4"}`)
	resolver, err := NewDimensionResolver(c)
	require.NoError(t, err)

	values, err := resolver.ResolveDimensions([]types.PricingDimension{
		{Key: "mode", Source: "derived.input_mode", Type: "string"},
	})
	require.NoError(t, err)
	require.Equal(t, "v2v", values.Values["mode"])
}

func TestDimensionResolver_MissingRequiredValue(t *testing.T) {
	c := newPricingJSONContext(`{"resolution":"720P"}`)
	resolver, err := NewDimensionResolver(c)
	require.NoError(t, err)

	values, err := resolver.ResolveDimensions([]types.PricingDimension{
		{Key: "duration", Source: "request.duration", Type: "number"},
	})
	require.NoError(t, err)
	require.True(t, values.Missing["duration"])
	require.Nil(t, values.Values["duration"])
}

func TestDimensionResolver_OptionalDefault(t *testing.T) {
	c := newPricingJSONContext(`{"resolution":"720P"}`)
	resolver, err := NewDimensionResolver(c)
	require.NoError(t, err)

	values, err := resolver.ResolveDimensions([]types.PricingDimension{
		{Key: "duration", Source: "request.duration", Type: "number", Optional: true, Default: 5.0},
	})
	require.NoError(t, err)
	require.Equal(t, 5.0, values.Values["duration"])
	require.Empty(t, values.Missing)
}

func TestDimensionResolver_MultipartFileInputMode(t *testing.T) {
	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	require.NoError(t, writer.WriteField("resolution", "720P"))
	part, err := writer.CreateFormFile("image", "frame.png")
	require.NoError(t, err)
	_, err = part.Write([]byte("fake"))
	require.NoError(t, err)
	require.NoError(t, writer.Close())

	c := newPricingContext(body.String(), writer.FormDataContentType())
	resolver, err := NewDimensionResolver(c)
	require.NoError(t, err)

	values, err := resolver.ResolveDimensions([]types.PricingDimension{
		{Key: "resolution", Source: "request.resolution", Type: "string"},
		{Key: "mode", Source: "derived.input_mode", Type: "string"},
	})
	require.NoError(t, err)
	require.Equal(t, "720P", values.Values["resolution"])
	require.Equal(t, "i2v", values.Values["mode"])
}

func TestDimensionResolver_UnsupportedContentType(t *testing.T) {
	c := newPricingContext("resolution=720P", "application/x-www-form-urlencoded")
	_, err := NewDimensionResolver(c)
	require.ErrorContains(t, err, "unsupported content-type")
}

func TestDimensionResolver_ResponseSourceUnsupported(t *testing.T) {
	c := newPricingJSONContext(`{"resolution":"720P"}`)
	resolver, err := NewDimensionResolver(c)
	require.NoError(t, err)

	_, err = resolver.ResolveDimensions([]types.PricingDimension{
		{Key: "tokens", Source: "response.usage.total_tokens", Type: "number"},
	})
	require.ErrorContains(t, err, "unsupported in schema v1")
}

func newPricingJSONContext(body string) *gin.Context {
	return newPricingContext(body, "application/json")
}

func newPricingContext(body string, contentType string) *gin.Context {
	gin.SetMode(gin.TestMode)
	w := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(w)
	req, _ := http.NewRequest(http.MethodPost, "/v1/video", strings.NewReader(body))
	req.Header.Set("Content-Type", contentType)
	c.Request = req
	return c
}