internal/ingestion/pipeline/cleanup_test.go
Rama main · réplica de lectura del repositorio original.
//
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
package pipeline
import (
"context"
"errors"
"sort"
"strings"
"testing"
"time"
"ragflow/internal/agent/runtime"
"ragflow/internal/ingestion/chunkcache"
"ragflow/internal/ingestion/component/globals"
"ragflow/internal/ingestion/mediacache"
)
type cleanupFakeStore struct {
kv map[string]string
sets map[string]map[string]bool
ttl map[string]time.Duration
failRead map[string]bool
bytes map[string]int64
}
func newCleanupFakeStore() *cleanupFakeStore {
return &cleanupFakeStore{
kv: map[string]string{},
sets: map[string]map[string]bool{},
ttl: map[string]time.Duration{},
failRead: map[string]bool{},
bytes: map[string]int64{},
}
}
func (f *cleanupFakeStore) Get(_ context.Context, key string) (string, error) {
return f.kv[key], nil
}
func (f *cleanupFakeStore) Set(_ context.Context, key, value string, exp time.Duration) bool {
f.kv[key] = value
f.ttl[key] = exp
return true
}
func (f *cleanupFakeStore) MediaPut(
_ context.Context,
mediaKey, digestKey, stateKey, digest, value string,
manifestKey, manifestMember string,
maxImages int,
maxBytes int64,
ttl time.Duration,
) (bool, bool, error) {
if f.sets[digestKey][digest] {
if f.kv[mediaKey] == "" {
f.kv[mediaKey] = value
f.ttl[mediaKey] = ttl
f.ttl[digestKey] = ttl
f.ttl[stateKey] = ttl
if manifestKey != "" {
if f.sets[manifestKey] == nil {
f.sets[manifestKey] = map[string]bool{}
}
f.sets[manifestKey][manifestMember] = true
f.ttl[manifestKey] = ttl
}
}
return true, false, nil
}
if maxImages > 0 && len(f.sets[digestKey]) >= maxImages {
return false, false, nil
}
if maxBytes > 0 && int64(len(value)) > maxBytes-f.bytes[stateKey] {
return false, false, nil
}
f.kv[mediaKey] = value
f.ttl[mediaKey] = ttl
if manifestKey != "" {
if f.sets[manifestKey] == nil {
f.sets[manifestKey] = map[string]bool{}
}
f.sets[manifestKey][manifestMember] = true
f.ttl[manifestKey] = ttl
}
if f.sets[digestKey] == nil {
f.sets[digestKey] = map[string]bool{}
}
f.sets[digestKey][digest] = true
f.bytes[stateKey] += int64(len(value))
f.ttl[digestKey] = ttl
f.ttl[stateKey] = ttl
return true, true, nil
}
func (f *cleanupFakeStore) SAdd(_ context.Context, key, member string) bool {
if f.sets[key] == nil {
f.sets[key] = map[string]bool{}
}
f.sets[key][member] = true
return true
}
func (f *cleanupFakeStore) SMembers(_ context.Context, key string) ([]string, error) {
if f.failRead[key] {
return nil, errors.New("kvrocks unavailable")
}
out := make([]string, 0, len(f.sets[key]))
for member := range f.sets[key] {
out = append(out, member)
}
sort.Strings(out)
return out, nil
}
func (f *cleanupFakeStore) Delete(_ context.Context, key string) bool {
delete(f.kv, key)
delete(f.sets, key)
delete(f.bytes, key)
return true
}
func (f *cleanupFakeStore) DeleteMany(_ context.Context, keys ...string) bool {
for _, k := range keys {
delete(f.kv, k)
delete(f.sets, k)
delete(f.bytes, k)
}
return true
}
func (f *cleanupFakeStore) Expire(_ context.Context, key string, exp time.Duration) bool {
f.ttl[key] = exp
return true
}
func cleanupTaskContext(taskID string) context.Context {
ctx := runtime.WithState(context.Background(), runtime.NewCanvasState("", ""))
globals.SetTaskID(ctx, taskID)
return ctx
}
func seedTaskManifests(t *testing.T, backend *cleanupFakeStore, taskID string) {
t.Helper()
ctx := cleanupTaskContext(taskID)
chunkcache.Set(ctx, backend, "kc:emb:chunk-1", "model result")
if _, err := mediacache.New(backend, taskID).Put(ctx, "html", "image/png", []byte("hello")); err != nil {
t.Fatalf("seed media cache: %v", err)
}
if len(backend.kv) != 2 || len(backend.sets) != 3 {
t.Fatalf("seeded state = %d values, %d sets, want 2 and 3", len(backend.kv), len(backend.sets))
}
}
func TestPurgeTaskStateWithBackendRemovesChunkAndMediaManifests(t *testing.T) {
backend := newCleanupFakeStore()
ctx := cleanupTaskContext("task-1")
seedTaskManifests(t, backend, "task-1")
if err := purgeTaskStateWithBackend(ctx, backend, "task-1"); err != nil {
t.Fatalf("purgeTaskStateWithBackend: %v", err)
}
if len(backend.kv) != 0 || len(backend.sets) != 0 {
t.Fatalf("state after purge = %v values, %v manifests, want empty", backend.kv, backend.sets)
}
}
func TestPurgeTaskStateWithBackendJoinsErrorsAndStillPurgesMedia(t *testing.T) {
backend := newCleanupFakeStore()
ctx := cleanupTaskContext("task-1")
seedTaskManifests(t, backend, "task-1")
backend.failRead["kc:manifest:task-1"] = true
err := purgeTaskStateWithBackend(ctx, backend, "task-1")
if err == nil || !strings.Contains(err.Error(), "chunk cache manifest") {
t.Fatalf("error = %v, want the chunk manifest failure", err)
}
for key := range backend.kv {
if strings.HasPrefix(key, "mc:") {
t.Fatalf("media value %s survived a failed chunk purge", key)
}
}
for key := range backend.sets {
if strings.HasPrefix(key, "mc:") {
t.Fatalf("media manifest %s survived a failed chunk purge", key)
}
}
if !backend.sets["kc:manifest:task-1"]["kc:emb:chunk-1"] {
t.Fatal("failed chunk purge lost the chunk manifest it must retry")
}
}
func TestPurgeTaskStateWithBackendRejectsEmptyTaskID(t *testing.T) {
backend := newCleanupFakeStore()
if err := purgeTaskStateWithBackend(context.Background(), backend, ""); err == nil {
t.Fatal("purge accepted an empty task id")
}
if len(backend.kv) != 0 || len(backend.sets) != 0 {
t.Fatal("empty task id touched backend state")
}
}