Function TestClientRank

modules/reranker-contextualai/clients/ranker_test.go:37–87

(t *testing.T)

)

func TestClientRank(t *testing.T) {
	testServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		assert.Equal(t, http.MethodPost, r.Method)
		assert.Equal(t, rerankPath, r.URL.Path)
		assert.Equal(t, applicationJSON, r.Header.Get(contentTypeHeader))
		assert.Contains(t, r.Header.Get("Authorization"), "Bearer")

		var req RankInput
		err := json.NewDecoder(r.Body).Decode(&req)
		require.NoError(t, err)

		assert.Equal(t, testQuery, req.Query)
		assert.Equal(t, []string{"doc1", "doc2", "doc3"}, req.Documents)
		assert.Equal(t, testModel, req.Model)

		response := RankResponse{
			Results: []RerankedResult{
				{Index: 0, RelevanceScore: 0.9},
				{Index: 2, RelevanceScore: 0.8},
				{Index: 1, RelevanceScore: 0.7},
			},
		}

		w.Header().Set(contentTypeHeader, applicationJSON)
		json.NewEncoder(w).Encode(response)
	}))
	defer testServer.Close()

	client := &client{
		apiKey:       testAPIKey,
		httpClient:   &http.Client{Timeout: 30 * time.Second},
		host:         testServer.URL,
		path:         rerankPath,
		maxDocuments: 1000,
		logger:       logrus.New(),
	}

	fakeConfig := fakeClassConfig{
		classConfig: map[string]any{
			"model": testModel,
		},
	}

	result, err := client.Rank(context.Background(), testQuery, []string{"doc1", "doc2", "doc3"}, fakeConfig)

	require.NoError(t, err)
	assert.Equal(t, testQuery, result.Query)
	assert.Len(t, result.DocumentScores, 3)
	assert.Equal(t, "doc1", result.DocumentScores[0].Document)
	assert.Equal(t, 0.9, result.DocumentScores[0].Score)
}

func TestClientMetaInfo(t *testing.T) {
	client := &client{}