79 lines
3.2 KiB
Go
79 lines
3.2 KiB
Go
|
package resumable
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"net/http/httptest"
|
||
|
|
"strings"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Tencent/WeKnora/internal/types"
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
func TestUploadSummaryKeepsProgressAndScopeWithoutSourcePayload(t *testing.T) {
|
||
|
|
s := uploadService(t)
|
||
|
|
secretSource := strings.Repeat("private-source-body-", 5000)
|
||
|
|
result, err := json.Marshal(map[string]any{"selection": "sampled", "source": map[string]string{"name": "source.txt", "text": secretSource}, "knowledge_ids": []string{"private-document"}})
|
||
|
|
require.NoError(t, err)
|
||
|
|
for _, fixture := range []struct {
|
||
|
|
id, owner, target string
|
||
|
|
tenant uint64
|
||
|
|
}{
|
||
|
|
{"owned", "alice", "allowed", 1}, {"other-user", "bob", "allowed", 1},
|
||
|
|
{"other-tenant", "alice", "allowed", 2}, {"revoked", "alice", "revoked", 1},
|
||
|
|
} {
|
||
|
|
require.NoError(t, s.db.Create(&Upload{ID: fixture.id, TenantID: fixture.tenant, OwnerID: fixture.owner,
|
||
|
|
Kind: "digital-human", TargetID: fixture.target, Name: fixture.id + ".txt", Size: 100,
|
||
|
|
Offset: 100, IngestOffset: 100, State: "completed", Fingerprint: "fingerprint",
|
||
|
|
Options: types.JSON(`{"purpose":"distillation"}`), Result: result, ExpiresAt: time.Now().Add(time.Hour)}).Error)
|
||
|
|
}
|
||
|
|
s.authorize = func(_ context.Context, _, target string) error {
|
||
|
|
if target == "revoked" {
|
||
|
|
return ErrForbidden
|
||
|
|
}
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
router := gin.New()
|
||
|
|
group := router.Group("/api/v1", func(c *gin.Context) { c.Request = c.Request.WithContext(uploadContext(1, "alice")) })
|
||
|
|
s.Register(group)
|
||
|
|
request := func(route string) []byte {
|
||
|
|
response := httptest.NewRecorder()
|
||
|
|
router.ServeHTTP(response, httptest.NewRequest("GET", route, nil))
|
||
|
|
require.Equal(t, 200, response.Code)
|
||
|
|
return response.Body.Bytes()
|
||
|
|
}
|
||
|
|
full := request("/api/v1/uploads")
|
||
|
|
summary := request("/api/v1/uploads?view=summary")
|
||
|
|
var envelope struct {
|
||
|
|
Data []Upload `json:"data"`
|
||
|
|
}
|
||
|
|
require.NoError(t, json.Unmarshal(summary, &envelope))
|
||
|
|
require.Len(t, envelope.Data, 1)
|
||
|
|
row := envelope.Data[0]
|
||
|
|
require.Equal(t, "owned", row.ID)
|
||
|
|
require.Equal(t, "completed", row.State)
|
||
|
|
require.EqualValues(t, 100, row.Offset)
|
||
|
|
require.EqualValues(t, 100, row.IngestOffset)
|
||
|
|
require.Equal(t, "fingerprint", row.Fingerprint)
|
||
|
|
require.JSONEq(t, `{"purpose":"distillation"}`, string(row.Options))
|
||
|
|
require.JSONEq(t, `{"selection":"sampled"}`, string(row.Result))
|
||
|
|
require.NotContains(t, string(summary), "private-source-body")
|
||
|
|
require.NotContains(t, string(summary), "private-document")
|
||
|
|
require.Less(t, len(summary)*10, len(full))
|
||
|
|
require.Contains(t, string(full), secretSource)
|
||
|
|
require.Contains(t, string(request("/api/v1/uploads/owned")), secretSource)
|
||
|
|
stored, err := s.Get(uploadContext(1, "alice"), "owned")
|
||
|
|
require.NoError(t, err)
|
||
|
|
require.JSONEq(t, string(result), string(stored.Result))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestUploadResultSummaryOmitsUnknownOrInvalidPayload(t *testing.T) {
|
||
|
|
for _, raw := range []string{`{}`, `null`, `[]`, `{"selection":"unexpected","source":{"text":"private"}}`, `{"selection":3}`, `invalid`} {
|
||
|
|
require.Nil(t, uploadResultSummary(types.JSON(raw)), raw)
|
||
|
|
}
|
||
|
|
require.JSONEq(t, `{"selection":"full"}`, string(uploadResultSummary(types.JSON(`{"selection":"full","source":{"text":"private"}}`))))
|
||
|
|
}
|