Files

48 lines
2.1 KiB
Go

package controller
import (
"context"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
fingerprint "github.com/QuantumNous/new-api/pkg/model_fingerprint"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
func TestAstraSamplingUsesIsolatedResponsesWithoutGuard(t *testing.T) {
requests := 0
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
requests++
bodyPreview, readErr := io.ReadAll(r.Body)
require.NoError(t, readErr)
r.Body = io.NopCloser(strings.NewReader(string(bodyPreview)))
assert.EqualValues(t, fingerprint.OutputTokens, gjson.GetBytes(bodyPreview, "max_output_tokens").Int())
assert.Equal(t, "low", gjson.GetBytes(bodyPreview, "reasoning.effort").String())
assert.Contains(t, gjson.GetBytes(bodyPreview, "input.0.content").String(), "355")
body, _ := io.ReadAll(r.Body)
if r.URL.Path != "/v1/responses" || gjson.GetBytes(body, "model").String() != fingerprint.Model || !gjson.GetBytes(body, "stream").Bool() {
t.Errorf("wrong probe request %s %s", r.URL.Path, body)
}
if gjson.GetBytes(body, "previous_response_id").Exists() || gjson.GetBytes(body, "conversation").Exists() {
t.Error("probe carried customer session")
}
w.Header().Set("Content-Type", "text/event-stream")
_, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"status\":\"completed\",\"model\":\"gpt-5.6-luna\",\"output\":[{\"type\":\"message\",\"content\":[{\"type\":\"output_text\",\"text\":\"[1,2,3]\"}]}]}}\n\n")
}))
defer server.Close()
url := server.URL
channel := &model.Channel{Id: 1, Type: 1, Key: "test-only", BaseURL: &url, Status: common.ChannelStatusEnabled, Models: fingerprint.Model}
r, count, err := runAstraFingerprintProbe(context.Background(), channel, 0)
if err != nil || r.DeclaredModel != "gpt-5.6-luna" || count < 292 || count > 304 || requests != 1 {
t.Fatalf("probe %+v count=%d requests=%d err=%v", r, count, requests, err)
}
}