mirror of
https://github.com/milvus-io/milvus.git
synced 2026-07-21 10:15:43 +00:00
986dcb42e6
## Summary - Bump **Go to 1.26.4** across CI workflows, Docker builders, RPM setup, and the `go` directive in all 4 modules (+ examples). - Refresh Go 1.26-compatible deps to latest: `sonic` v1.15.2, `mockey` v1.4.6, `gopkg` v0.1.4, `sonic/loader` v0.5.1, `cpuid/v2` v2.3.0. sonic 1.14.x fails to build on the 1.26 compiler (bytedance/sonic#895). - **datacoord tests:** replace the anonymous `struct{ Iface }` mockey idiom with **named helper types** (`embeddedHandler`/`embeddedBroadcastAPI`/`embeddedBroker`/`embeddedAllocator`). ## Why the datacoord test change Go 1.26's `printf` vet pass (run automatically by `go test`) calls `x/tools/refactor/satisfy.Finder` unconditionally and **panics `(*ast.StructType)`** on a **method expression of `*struct{ Iface }`** — an upstream Go toolchain bug. Only `internal/datacoord` used this mockey idiom (31 sites in 2 test files; verified repo-wide it's the only place), which is why only `build-ut-cov`/`ut-go` failed on `internal/datacoord [build failed]`. Switching to a **named type** makes the method-expression receiver an identifier instead of a struct literal, so `satisfy.Finder` no longer panics. mockey usage and test behavior are unchanged. Minimal repro: `type I interface{ Logf(string, ...any) }` + `var _ = (*struct{ I }).Logf` under `go 1.26` → `go vet` panics; named type does not. ## Notes - `.env` builder-image tag intentionally not hand-edited (auto-bumped post-merge by `bump-builder-version`). - The bug also warrants an upstream report to golang/go (reproducer above). ## Test plan - [x] `go mod tidy` clean across all four modules - [x] sonic 1.15.2 + mockey compile on go 1.26.4 - [x] reproduced the vet panic; verified named-type method expression avoids it; gofmt-clean - [ ] Full CI (build / build-ut-cov / ut-go / ut-cpp / integration) 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Signed-off-by: xiaofanluan <xiaofan.luan@zilliz.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1096 lines
36 KiB
Go
1096 lines
36 KiB
Go
// Licensed to the LF AI & Data foundation under one
|
|
// or more contributor license agreements. See the NOTICE file
|
|
// distributed with this work for additional information
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
// to you 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 datacoord
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/bytedance/mockey"
|
|
"github.com/cockroachdb/errors"
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/datapb"
|
|
"github.com/milvus-io/milvus/pkg/v3/proto/indexpb"
|
|
)
|
|
|
|
// ==================== Helper Functions for Meta Tests ====================
|
|
|
|
func createMetaTestRefreshMeta(t *testing.T, jobs []*datapb.ExternalCollectionRefreshJob, tasks []*datapb.ExternalCollectionRefreshTask) *externalCollectionRefreshMeta {
|
|
catalog := &stubCatalog{
|
|
jobs: jobs,
|
|
tasks: tasks,
|
|
}
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
return meta
|
|
}
|
|
|
|
// ==================== Test Functions ====================
|
|
|
|
func TestExternalCollectionRefreshMeta_NewMeta(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, CollectionName: "test_collection"},
|
|
}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, CollectionId: 100},
|
|
}
|
|
catalog := &stubCatalog{jobs: jobs, tasks: tasks}
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, meta)
|
|
|
|
// Verify job loaded
|
|
job := meta.GetJob(1)
|
|
assert.NotNil(t, job)
|
|
assert.Equal(t, int64(1), job.GetJobId())
|
|
assert.Equal(t, int64(100), job.GetCollectionId())
|
|
|
|
// Verify task loaded
|
|
task := meta.GetTask(1001)
|
|
assert.NotNil(t, task)
|
|
assert.Equal(t, int64(1001), task.GetTaskId())
|
|
})
|
|
|
|
t.Run("list_jobs_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
|
|
// Mock ListExternalCollectionRefreshJobs to return error
|
|
mockList := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, errors.New("list jobs error")).Build()
|
|
defer mockList.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, meta)
|
|
})
|
|
|
|
t.Run("list_tasks_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
|
|
// Mock ListExternalCollectionRefreshTasks to return error
|
|
mockList := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(nil, errors.New("list tasks error")).Build()
|
|
defer mockList.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.Error(t, err)
|
|
assert.Nil(t, meta)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_AddJob(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
job := &datapb.ExternalCollectionRefreshJob{
|
|
JobId: 1,
|
|
CollectionId: 100,
|
|
CollectionName: "test_collection",
|
|
State: indexpb.JobState_JobStateInit,
|
|
StartTime: time.Now().UnixMilli(),
|
|
}
|
|
|
|
err := meta.AddJob(job)
|
|
assert.NoError(t, err)
|
|
|
|
// Verify job added
|
|
got := meta.GetJob(1)
|
|
assert.NotNil(t, got)
|
|
assert.Equal(t, int64(1), got.GetJobId())
|
|
assert.Equal(t, int64(100), got.GetCollectionId())
|
|
})
|
|
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.NoError(t, err)
|
|
|
|
// Mock SaveExternalCollectionRefreshJob to return error
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshJob).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
job := &datapb.ExternalCollectionRefreshJob{
|
|
JobId: 1,
|
|
CollectionId: 100,
|
|
}
|
|
|
|
err = meta.AddJob(job)
|
|
assert.Error(t, err)
|
|
|
|
// Verify job not added
|
|
got := meta.GetJob(1)
|
|
assert.Nil(t, got)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_GetJob(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, CollectionName: "test_collection"},
|
|
}
|
|
catalog := &stubCatalog{jobs: jobs}
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.NoError(t, err)
|
|
|
|
t.Run("job_exists", func(t *testing.T) {
|
|
job := meta.GetJob(1)
|
|
assert.NotNil(t, job)
|
|
assert.Equal(t, int64(1), job.GetJobId())
|
|
})
|
|
|
|
t.Run("job_not_exists", func(t *testing.T) {
|
|
job := meta.GetJob(999)
|
|
assert.Nil(t, job)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_GetActiveJobByCollectionID(t *testing.T) {
|
|
now := time.Now().UnixMilli()
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateInit, StartTime: now - 1000},
|
|
{JobId: 2, CollectionId: 100, State: indexpb.JobState_JobStateInProgress, StartTime: now},
|
|
{JobId: 3, CollectionId: 100, State: indexpb.JobState_JobStateFinished, StartTime: now + 1000},
|
|
{JobId: 4, CollectionId: 200, State: indexpb.JobState_JobStateInit, StartTime: now},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
t.Run("has_active_job", func(t *testing.T) {
|
|
job := meta.GetActiveJobByCollectionID(100)
|
|
assert.NotNil(t, job)
|
|
// Should return the newest active job (jobId=2 with InProgress state)
|
|
assert.Equal(t, int64(2), job.GetJobId())
|
|
})
|
|
|
|
t.Run("no_active_job", func(t *testing.T) {
|
|
// Collection 300 has no jobs
|
|
job := meta.GetActiveJobByCollectionID(300)
|
|
assert.Nil(t, job)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_ListJobsByCollectionID(t *testing.T) {
|
|
now := time.Now().UnixMilli()
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, StartTime: now - 2000},
|
|
{JobId: 2, CollectionId: 100, StartTime: now - 1000},
|
|
{JobId: 3, CollectionId: 100, StartTime: now},
|
|
{JobId: 4, CollectionId: 200, StartTime: now},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
t.Run("all_jobs", func(t *testing.T) {
|
|
jobs := meta.ListJobsByCollectionID(100)
|
|
assert.Len(t, jobs, 3)
|
|
// Should be sorted by StartTime descending
|
|
assert.Equal(t, int64(3), jobs[0].GetJobId())
|
|
assert.Equal(t, int64(2), jobs[1].GetJobId())
|
|
assert.Equal(t, int64(1), jobs[2].GetJobId())
|
|
})
|
|
|
|
t.Run("no_jobs", func(t *testing.T) {
|
|
jobs := meta.ListJobsByCollectionID(300)
|
|
assert.Nil(t, jobs)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_GetAllJobs(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, StartTime: 100},
|
|
{JobId: 2, CollectionId: 200, StartTime: 200},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
allJobs := meta.GetAllJobs()
|
|
assert.Len(t, allJobs, 2)
|
|
assert.NotNil(t, allJobs[1])
|
|
assert.NotNil(t, allJobs[2])
|
|
|
|
listedJobs := meta.ListAllJobs()
|
|
assert.Len(t, listedJobs, 2)
|
|
assert.Equal(t, int64(2), listedJobs[0].GetJobId())
|
|
assert.Equal(t, int64(1), listedJobs[1].GetJobId())
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_UpdateJobState(t *testing.T) {
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateInit},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(jobs, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(nil, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
// Mock save to fail
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshJob).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
applied, err := meta.UpdateJobState(1, indexpb.JobState_JobStateInProgress, "")
|
|
assert.Error(t, err)
|
|
assert.False(t, applied)
|
|
|
|
// State should remain Init
|
|
assert.Equal(t, indexpb.JobState_JobStateInit, meta.GetJob(1).GetState())
|
|
})
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateInit},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
applied, err := meta.UpdateJobState(1, indexpb.JobState_JobStateInProgress, "")
|
|
assert.NoError(t, err)
|
|
assert.True(t, applied)
|
|
|
|
job := meta.GetJob(1)
|
|
assert.Equal(t, indexpb.JobState_JobStateInProgress, job.GetState())
|
|
})
|
|
|
|
t.Run("finished_sets_end_time_and_progress", func(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
applied, err := meta.UpdateJobState(1, indexpb.JobState_JobStateFinished, "")
|
|
assert.NoError(t, err)
|
|
assert.True(t, applied)
|
|
|
|
job := meta.GetJob(1)
|
|
assert.Equal(t, indexpb.JobState_JobStateFinished, job.GetState())
|
|
assert.Equal(t, int64(100), job.GetProgress())
|
|
assert.Greater(t, job.GetEndTime(), int64(0))
|
|
})
|
|
|
|
t.Run("failed_sets_end_time_and_reason", func(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
applied, err := meta.UpdateJobState(1, indexpb.JobState_JobStateFailed, "timeout")
|
|
assert.NoError(t, err)
|
|
assert.True(t, applied)
|
|
|
|
job := meta.GetJob(1)
|
|
assert.Equal(t, indexpb.JobState_JobStateFailed, job.GetState())
|
|
assert.Equal(t, "timeout", job.GetFailReason())
|
|
assert.Greater(t, job.GetEndTime(), int64(0))
|
|
})
|
|
|
|
t.Run("terminal_state_guard_skips_write", func(t *testing.T) {
|
|
// Once a job is Finished, a follow-up UpdateJobState(Failed) must
|
|
// NOT persist the transition and MUST return applied=false so the
|
|
// caller can distinguish "silently skipped" from "persisted". This
|
|
// is the signal tryTimeoutJob relies on to avoid poisoning the
|
|
// manager's notifiedJobs dedup map during a race with the eager
|
|
// Finished path.
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateFinished, Progress: 100},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
applied, err := meta.UpdateJobState(1, indexpb.JobState_JobStateFailed, "timeout")
|
|
assert.NoError(t, err)
|
|
assert.False(t, applied, "terminal-state guard must report applied=false")
|
|
|
|
job := meta.GetJob(1)
|
|
assert.Equal(t, indexpb.JobState_JobStateFinished, job.GetState(), "state must remain Finished")
|
|
})
|
|
|
|
t.Run("job_not_found", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
applied, err := meta.UpdateJobState(999, indexpb.JobState_JobStateInProgress, "")
|
|
assert.Error(t, err)
|
|
assert.False(t, applied)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_UpdateJobStateWithPreApply(t *testing.T) {
|
|
t.Run("pre_apply_failure_marks_job_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(jobs, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(nil, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
var savedJob *datapb.ExternalCollectionRefreshJob
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshJob).
|
|
To(func(_ context.Context, job *datapb.ExternalCollectionRefreshJob) error {
|
|
savedJob = job
|
|
return nil
|
|
}).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
applied, err := meta.UpdateJobStateWithPreApply(
|
|
1,
|
|
indexpb.JobState_JobStateFinished,
|
|
"",
|
|
func(*datapb.ExternalCollectionRefreshJob) error {
|
|
return errors.New("apply failed")
|
|
})
|
|
|
|
assert.True(t, applied)
|
|
assert.Error(t, err)
|
|
assert.Contains(t, err.Error(), "apply failed")
|
|
assert.NotNil(t, savedJob)
|
|
assert.Equal(t, indexpb.JobState_JobStateFailed, savedJob.GetState())
|
|
assert.Equal(t, "apply failed", savedJob.GetFailReason())
|
|
assert.Equal(t, indexpb.JobState_JobStateFailed, meta.GetJob(1).GetState())
|
|
})
|
|
|
|
t.Run("pre_apply_success_save_job_failure_keeps_original_job", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(jobs, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(nil, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshJob).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
preApplyCalled := false
|
|
applied, err := meta.UpdateJobStateWithPreApply(
|
|
1,
|
|
indexpb.JobState_JobStateFinished,
|
|
"",
|
|
func(*datapb.ExternalCollectionRefreshJob) error {
|
|
preApplyCalled = true
|
|
return nil
|
|
})
|
|
|
|
assert.False(t, applied)
|
|
assert.Error(t, err)
|
|
assert.True(t, preApplyCalled)
|
|
assert.Equal(t, indexpb.JobState_JobStateInProgress, meta.GetJob(1).GetState())
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_UpdateJobProgress(t *testing.T) {
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, Progress: 0},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(jobs, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(nil, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshJob).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
err = meta.UpdateJobProgress(1, 50)
|
|
assert.Error(t, err)
|
|
assert.Equal(t, int64(0), meta.GetJob(1).GetProgress())
|
|
})
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, Progress: 0},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
err := meta.UpdateJobProgress(1, 50)
|
|
assert.NoError(t, err)
|
|
|
|
job := meta.GetJob(1)
|
|
assert.Equal(t, int64(50), job.GetProgress())
|
|
})
|
|
|
|
t.Run("job_not_found", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
err := meta.UpdateJobProgress(999, 50)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_AddTaskIDToJob(t *testing.T) {
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, TaskIds: []int64{}},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(jobs, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(nil, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshJob).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
err = meta.AddTaskIDToJob(1, 1001)
|
|
assert.Error(t, err)
|
|
})
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100, TaskIds: []int64{}},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, nil)
|
|
|
|
err := meta.AddTaskIDToJob(1, 1001)
|
|
assert.NoError(t, err)
|
|
|
|
job := meta.GetJob(1)
|
|
assert.Contains(t, job.GetTaskIds(), int64(1001))
|
|
})
|
|
|
|
t.Run("job_not_found", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
err := meta.AddTaskIDToJob(999, 1001)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_DropJob(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
t.Run("drop_task_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100},
|
|
}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, CollectionId: 100},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(jobs, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockDropTask := mockey.Mock((*stubCatalog).DropExternalCollectionRefreshTask).Return(errors.New("drop task error")).Build()
|
|
defer mockDropTask.UnPatch()
|
|
|
|
err = meta.DropJob(ctx, 1)
|
|
assert.Error(t, err)
|
|
|
|
// Job should still exist
|
|
assert.NotNil(t, meta.GetJob(1))
|
|
})
|
|
|
|
t.Run("drop_job_catalog_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(jobs, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(nil, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.NoError(t, err)
|
|
|
|
// Mock drop job to fail
|
|
mockDropJob := mockey.Mock((*stubCatalog).DropExternalCollectionRefreshJob).Return(errors.New("drop job error")).Build()
|
|
defer mockDropJob.UnPatch()
|
|
|
|
err = meta.DropJob(ctx, 1)
|
|
assert.Error(t, err)
|
|
|
|
// Job should still exist
|
|
assert.NotNil(t, meta.GetJob(1))
|
|
})
|
|
|
|
t.Run("success_with_tasks", func(t *testing.T) {
|
|
jobs := []*datapb.ExternalCollectionRefreshJob{
|
|
{JobId: 1, CollectionId: 100},
|
|
}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, CollectionId: 100},
|
|
{TaskId: 1002, JobId: 1, CollectionId: 100},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, jobs, tasks)
|
|
|
|
err := meta.DropJob(ctx, 1)
|
|
assert.NoError(t, err)
|
|
|
|
// Verify job and tasks removed
|
|
assert.Nil(t, meta.GetJob(1))
|
|
assert.Nil(t, meta.GetTask(1001))
|
|
assert.Nil(t, meta.GetTask(1002))
|
|
})
|
|
|
|
t.Run("job_not_exists", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
// Should not error if job doesn't exist
|
|
err := meta.DropJob(ctx, 999)
|
|
assert.NoError(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_AddTask(t *testing.T) {
|
|
t.Run("success", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
task := &datapb.ExternalCollectionRefreshTask{
|
|
TaskId: 1001,
|
|
JobId: 1,
|
|
CollectionId: 100,
|
|
State: indexpb.JobState_JobStateInit,
|
|
}
|
|
|
|
err := meta.AddTask(task)
|
|
assert.NoError(t, err)
|
|
|
|
// Verify task added
|
|
got := meta.GetTask(1001)
|
|
assert.NotNil(t, got)
|
|
assert.Equal(t, int64(1001), got.GetTaskId())
|
|
})
|
|
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
// Mock SaveExternalCollectionRefreshTask to return error
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshTask).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
task := &datapb.ExternalCollectionRefreshTask{
|
|
TaskId: 1001,
|
|
JobId: 1,
|
|
}
|
|
|
|
err = meta.AddTask(task)
|
|
assert.Error(t, err)
|
|
|
|
// Verify task not added
|
|
got := meta.GetTask(1001)
|
|
assert.Nil(t, got)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_GetTask(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, CollectionId: 100},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
t.Run("task_exists", func(t *testing.T) {
|
|
task := meta.GetTask(1001)
|
|
assert.NotNil(t, task)
|
|
assert.Equal(t, int64(1001), task.GetTaskId())
|
|
})
|
|
|
|
t.Run("task_not_exists", func(t *testing.T) {
|
|
task := meta.GetTask(9999)
|
|
assert.Nil(t, task)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_GetTasksByJobID(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, CollectionId: 100},
|
|
{TaskId: 1002, JobId: 1, CollectionId: 100},
|
|
{TaskId: 2001, JobId: 2, CollectionId: 200},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
t.Run("has_tasks", func(t *testing.T) {
|
|
tasks := meta.GetTasksByJobID(1)
|
|
assert.Len(t, tasks, 2)
|
|
})
|
|
|
|
t.Run("no_tasks", func(t *testing.T) {
|
|
tasks := meta.GetTasksByJobID(999)
|
|
assert.Nil(t, tasks)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_GetAllTasks(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1},
|
|
{TaskId: 1002, JobId: 1},
|
|
{TaskId: 2001, JobId: 2},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
allTasks := meta.GetAllTasks()
|
|
assert.Len(t, allTasks, 3)
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_GetTaskState(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
t.Run("task_exists", func(t *testing.T) {
|
|
state := meta.GetTaskState(1001)
|
|
assert.Equal(t, indexpb.JobState_JobStateInProgress, state)
|
|
})
|
|
|
|
t.Run("task_not_exists", func(t *testing.T) {
|
|
state := meta.GetTaskState(9999)
|
|
assert.Equal(t, indexpb.JobState_JobStateNone, state)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_UpdateTaskState(t *testing.T) {
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateInit},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshTask).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
err = meta.UpdateTaskState(1001, indexpb.JobState_JobStateInProgress, "")
|
|
assert.Error(t, err)
|
|
assert.Equal(t, indexpb.JobState_JobStateInit, meta.GetTask(1001).GetState())
|
|
})
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateInit},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
err := meta.UpdateTaskState(1001, indexpb.JobState_JobStateInProgress, "")
|
|
assert.NoError(t, err)
|
|
|
|
task := meta.GetTask(1001)
|
|
assert.Equal(t, indexpb.JobState_JobStateInProgress, task.GetState())
|
|
})
|
|
|
|
t.Run("failed_state", func(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
err := meta.UpdateTaskState(1001, indexpb.JobState_JobStateFailed, "connection timeout")
|
|
assert.NoError(t, err)
|
|
|
|
task := meta.GetTask(1001)
|
|
assert.Equal(t, indexpb.JobState_JobStateFailed, task.GetState())
|
|
assert.Equal(t, "connection timeout", task.GetFailReason())
|
|
})
|
|
|
|
t.Run("task_not_found", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
err := meta.UpdateTaskState(9999, indexpb.JobState_JobStateInProgress, "")
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_UpdateTaskResult(t *testing.T) {
|
|
t.Run("persists_result_and_clones_segments", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
var savedTask *datapb.ExternalCollectionRefreshTask
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshTask).
|
|
To(func(_ context.Context, task *datapb.ExternalCollectionRefreshTask) error {
|
|
savedTask = task
|
|
return nil
|
|
}).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
updatedSegment := &datapb.SegmentInfo{ID: 10, CollectionID: 100, NumOfRows: 7}
|
|
err = meta.UpdateTaskResult(
|
|
1001,
|
|
indexpb.JobState_JobStateFinished,
|
|
"",
|
|
[]int64{1, 2},
|
|
[]*datapb.SegmentInfo{updatedSegment},
|
|
)
|
|
assert.NoError(t, err)
|
|
assert.NotNil(t, savedTask)
|
|
assert.Equal(t, indexpb.JobState_JobStateFinished, savedTask.GetState())
|
|
assert.Equal(t, int64(100), savedTask.GetProgress())
|
|
assert.True(t, savedTask.GetResultReady())
|
|
assert.Equal(t, []int64{1, 2}, savedTask.GetKeptSegments())
|
|
assert.Len(t, savedTask.GetUpdatedSegments(), 1)
|
|
assert.Equal(t, int64(10), savedTask.GetUpdatedSegments()[0].GetID())
|
|
assert.Equal(t, int64(7), savedTask.GetUpdatedSegments()[0].GetNumOfRows())
|
|
|
|
updatedSegment.NumOfRows = 99
|
|
task := meta.GetTask(1001)
|
|
assert.Equal(t, int64(7), task.GetUpdatedSegments()[0].GetNumOfRows())
|
|
assert.Equal(t, int64(7), savedTask.GetUpdatedSegments()[0].GetNumOfRows())
|
|
})
|
|
|
|
t.Run("save_failed_keeps_original_task", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateInProgress},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshTask).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
err = meta.UpdateTaskResult(
|
|
1001,
|
|
indexpb.JobState_JobStateFinished,
|
|
"",
|
|
[]int64{1},
|
|
[]*datapb.SegmentInfo{{ID: 10, CollectionID: 100, NumOfRows: 7}},
|
|
)
|
|
assert.Error(t, err)
|
|
|
|
task := meta.GetTask(1001)
|
|
assert.Equal(t, indexpb.JobState_JobStateInProgress, task.GetState())
|
|
assert.Empty(t, task.GetKeptSegments())
|
|
assert.Empty(t, task.GetUpdatedSegments())
|
|
})
|
|
|
|
t.Run("task_not_found", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
err := meta.UpdateTaskResult(
|
|
9999,
|
|
indexpb.JobState_JobStateFinished,
|
|
"",
|
|
[]int64{1},
|
|
[]*datapb.SegmentInfo{{ID: 10}},
|
|
)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_ClearTaskResultsByJobID_PartialFailure(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{
|
|
TaskId: 1001,
|
|
JobId: 1,
|
|
State: indexpb.JobState_JobStateFinished,
|
|
ResultReady: true,
|
|
KeptSegments: []int64{1},
|
|
UpdatedSegments: []*datapb.SegmentInfo{{ID: 10}},
|
|
},
|
|
{
|
|
TaskId: 1002,
|
|
JobId: 1,
|
|
State: indexpb.JobState_JobStateFinished,
|
|
ResultReady: true,
|
|
KeptSegments: []int64{2},
|
|
UpdatedSegments: []*datapb.SegmentInfo{{ID: 20}},
|
|
},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
saveCalls := 0
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshTask).
|
|
To(func(_ context.Context, task *datapb.ExternalCollectionRefreshTask) error {
|
|
saveCalls++
|
|
if task.GetTaskId() == 1002 {
|
|
return errors.New("save error")
|
|
}
|
|
return nil
|
|
}).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
err = meta.ClearTaskResultsByJobID(1)
|
|
assert.Error(t, err)
|
|
assert.Equal(t, 2, saveCalls)
|
|
assert.Empty(t, meta.GetTask(1001).GetKeptSegments())
|
|
assert.Empty(t, meta.GetTask(1001).GetUpdatedSegments())
|
|
assert.Equal(t, []int64{2}, meta.GetTask(1002).GetKeptSegments())
|
|
assert.Len(t, meta.GetTask(1002).GetUpdatedSegments(), 1)
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_UpdateTaskProgress(t *testing.T) {
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, Progress: 0},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshTask).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
err = meta.UpdateTaskProgress(1001, 50)
|
|
assert.Error(t, err)
|
|
assert.Equal(t, int64(0), meta.GetTask(1001).GetProgress())
|
|
})
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, Progress: 0},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
err := meta.UpdateTaskProgress(1001, 50)
|
|
assert.NoError(t, err)
|
|
|
|
task := meta.GetTask(1001)
|
|
assert.Equal(t, int64(50), task.GetProgress())
|
|
})
|
|
|
|
t.Run("task_not_found", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
err := meta.UpdateTaskProgress(9999, 50)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_UpdateTaskVersion(t *testing.T) {
|
|
t.Run("save_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, Version: 0, NodeId: 0},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(context.Background(), catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockSave := mockey.Mock((*stubCatalog).SaveExternalCollectionRefreshTask).Return(errors.New("save error")).Build()
|
|
defer mockSave.UnPatch()
|
|
|
|
err = meta.UpdateTaskVersion(1001, 10)
|
|
assert.Error(t, err)
|
|
// Version and NodeId should remain unchanged
|
|
task := meta.GetTask(1001)
|
|
assert.Equal(t, int64(0), task.GetVersion())
|
|
assert.Equal(t, int64(0), task.GetNodeId())
|
|
})
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, Version: 0, NodeId: 0},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
err := meta.UpdateTaskVersion(1001, 10)
|
|
assert.NoError(t, err)
|
|
|
|
task := meta.GetTask(1001)
|
|
assert.Equal(t, int64(1), task.GetVersion())
|
|
assert.Equal(t, int64(10), task.GetNodeId())
|
|
})
|
|
|
|
t.Run("task_not_found", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
err := meta.UpdateTaskVersion(9999, 10)
|
|
assert.Error(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_DropTask(t *testing.T) {
|
|
ctx := context.Background()
|
|
|
|
t.Run("drop_failed", func(t *testing.T) {
|
|
catalog := &stubCatalog{}
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1},
|
|
}
|
|
mockListJobs := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshJobs).Return(nil, nil).Build()
|
|
defer mockListJobs.UnPatch()
|
|
mockListTasks := mockey.Mock((*stubCatalog).ListExternalCollectionRefreshTasks).Return(tasks, nil).Build()
|
|
defer mockListTasks.UnPatch()
|
|
|
|
meta, err := newExternalCollectionRefreshMeta(ctx, catalog)
|
|
assert.NoError(t, err)
|
|
|
|
mockDrop := mockey.Mock((*stubCatalog).DropExternalCollectionRefreshTask).Return(errors.New("drop error")).Build()
|
|
defer mockDrop.UnPatch()
|
|
|
|
err = meta.DropTask(ctx, 1001)
|
|
assert.Error(t, err)
|
|
|
|
// Task should still exist
|
|
assert.NotNil(t, meta.GetTask(1001))
|
|
})
|
|
|
|
t.Run("success", func(t *testing.T) {
|
|
tasks := []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1},
|
|
}
|
|
meta := createMetaTestRefreshMeta(t, nil, tasks)
|
|
|
|
err := meta.DropTask(ctx, 1001)
|
|
assert.NoError(t, err)
|
|
|
|
// Verify task removed
|
|
assert.Nil(t, meta.GetTask(1001))
|
|
})
|
|
|
|
t.Run("task_not_exists", func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, nil)
|
|
|
|
// Should not error if task doesn't exist
|
|
err := meta.DropTask(ctx, 9999)
|
|
assert.NoError(t, err)
|
|
})
|
|
}
|
|
|
|
func TestExternalCollectionRefreshMeta_AggregateJobStateFromTasks(t *testing.T) {
|
|
testCases := []struct {
|
|
name string
|
|
tasks []*datapb.ExternalCollectionRefreshTask
|
|
expectedState indexpb.JobState
|
|
expectedProgress int64
|
|
}{
|
|
{
|
|
name: "no_tasks",
|
|
tasks: nil,
|
|
expectedState: indexpb.JobState_JobStateNone,
|
|
expectedProgress: 0,
|
|
},
|
|
{
|
|
name: "all_init",
|
|
tasks: []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateInit, Progress: 0},
|
|
{TaskId: 1002, JobId: 1, State: indexpb.JobState_JobStateInit, Progress: 0},
|
|
},
|
|
expectedState: indexpb.JobState_JobStateInit,
|
|
expectedProgress: 0,
|
|
},
|
|
{
|
|
name: "all_finished",
|
|
tasks: []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateFinished, Progress: 100},
|
|
{TaskId: 1002, JobId: 1, State: indexpb.JobState_JobStateFinished, Progress: 100},
|
|
},
|
|
expectedState: indexpb.JobState_JobStateFinished,
|
|
expectedProgress: 100,
|
|
},
|
|
{
|
|
name: "has_in_progress",
|
|
tasks: []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateFinished, Progress: 100},
|
|
{TaskId: 1002, JobId: 1, State: indexpb.JobState_JobStateInProgress, Progress: 50},
|
|
},
|
|
expectedState: indexpb.JobState_JobStateInProgress,
|
|
expectedProgress: 75, // (100+50)/2
|
|
},
|
|
{
|
|
name: "has_failed",
|
|
tasks: []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateFinished, Progress: 100},
|
|
{TaskId: 1002, JobId: 1, State: indexpb.JobState_JobStateInProgress, Progress: 50},
|
|
{TaskId: 1003, JobId: 1, State: indexpb.JobState_JobStateFailed, Progress: 30},
|
|
},
|
|
expectedState: indexpb.JobState_JobStateFailed,
|
|
expectedProgress: 60, // (100+50+30)/3
|
|
},
|
|
{
|
|
name: "has_retry",
|
|
tasks: []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateFinished, Progress: 100},
|
|
{TaskId: 1002, JobId: 1, State: indexpb.JobState_JobStateRetry, Progress: 20},
|
|
},
|
|
expectedState: indexpb.JobState_JobStateRetry,
|
|
expectedProgress: 60, // (100+20)/2
|
|
},
|
|
{
|
|
name: "init_over_finished",
|
|
tasks: []*datapb.ExternalCollectionRefreshTask{
|
|
{TaskId: 1001, JobId: 1, State: indexpb.JobState_JobStateFinished, Progress: 100},
|
|
{TaskId: 1002, JobId: 1, State: indexpb.JobState_JobStateInit, Progress: 0},
|
|
},
|
|
expectedState: indexpb.JobState_JobStateInit,
|
|
expectedProgress: 50, // (100+0)/2
|
|
},
|
|
}
|
|
|
|
for _, tc := range testCases {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
meta := createMetaTestRefreshMeta(t, nil, tc.tasks)
|
|
|
|
state, progress := meta.AggregateJobStateFromTasks(1)
|
|
assert.Equal(t, tc.expectedState, state)
|
|
assert.Equal(t, tc.expectedProgress, progress)
|
|
})
|
|
}
|
|
}
|