Files
wehub-resource-sync 498b235461
Build and test / Build and test AMD64 Ubuntu 22.04 (push) Failing after 0s
Publish Builder / amazonlinux2023 (push) Failing after 1s
Build and test / UT for Go (push) Has been skipped
Publish KRTE Images / KRTE (push) Failing after 1s
Build and test / Integration Test (push) Has been skipped
Build and test / Upload Code Coverage (push) Has been skipped
Publish Builder / rockylinux9 (push) Failing after 1s
Publish Builder / ubuntu22.04 (push) Failing after 0s
Publish Builder / ubuntu24.04 (push) Failing after 0s
Publish Gpu Builder / publish-gpu-builder (push) Failing after 1s
Publish Test Images / PyTest (push) Failing after 0s
Build and test / UT for Cpp (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:31:17 +08:00

1431 lines
54 KiB
Go

package service
import (
"context"
"io"
"testing"
"time"
"github.com/bytedance/mockey"
"github.com/cockroachdb/errors"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/milvus-io/milvus-proto/go-api/v3/commonpb"
"github.com/milvus-io/milvus/internal/coordinator/snmanager"
"github.com/milvus-io/milvus/internal/distributed/streaming"
"github.com/milvus-io/milvus/internal/mocks/distributed/mock_streaming"
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_balancer"
"github.com/milvus-io/milvus/internal/mocks/streamingcoord/server/mock_broadcaster"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/balance"
"github.com/milvus-io/milvus/internal/streamingcoord/server/balancer/channel"
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster"
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/broadcast"
"github.com/milvus-io/milvus/internal/streamingcoord/server/broadcaster/registry"
"github.com/milvus-io/milvus/internal/streamingcoord/server/resource"
"github.com/milvus-io/milvus/pkg/v3/mocks/proto/mock_streamingpb"
"github.com/milvus-io/milvus/pkg/v3/proto/streamingpb"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/message"
"github.com/milvus-io/milvus/pkg/v3/streaming/util/types"
"github.com/milvus-io/milvus/pkg/v3/util/syncutil"
)
func TestAssignmentService(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
// Set up the balancer
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
})
b.EXPECT().Close().Return().Maybe()
balance.Register(b)
// Mock channel.GetClusterChannels to avoid blocking on unregistered singleton.
mockGetClusterChannels := mockey.Mock(channel.GetClusterChannels).Return(message.ClusterChannels{
Channels: []string{"by-dev-1"},
ControlChannel: "by-dev-1_vcchan",
}).Build()
defer mockGetClusterChannels.UnPatch()
// Set up the broadcaster
fb := syncutil.NewFuture[broadcaster.Broadcaster]()
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mb := mock_broadcaster.NewMockBroadcaster(t)
fb.Set(mb)
mba.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(&types.BroadcastAppendResult{}, nil).Maybe()
mba.EXPECT().Close().Return().Maybe()
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Ack(mock.Anything, mock.Anything).Return(nil).Maybe()
mb.EXPECT().LegacyAck(mock.Anything, mock.Anything, mock.Anything).Return(nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
// Test assignment discover
as := NewAssignmentService()
ss := mock_streamingpb.NewMockStreamingCoordAssignmentService_AssignmentDiscoverServer(t)
ss.EXPECT().Context().Return(context.Background()).Maybe()
ss.EXPECT().Recv().Return(nil, io.EOF).Maybe()
ss.EXPECT().Send(mock.Anything).Return(io.EOF).Maybe()
err := as.AssignmentDiscover(ss)
assert.Error(t, err)
// Test update WAL balance policy
b.EXPECT().UpdateBalancePolicy(context.Background(), mock.Anything).Return(&streamingpb.UpdateWALBalancePolicyResponse{}, nil).Maybe()
as.UpdateWALBalancePolicy(context.Background(), &streamingpb.UpdateWALBalancePolicyRequest{})
// Test update replicate configuration
// Test illegal replicate configuration
cfg := &commonpb.ReplicateConfiguration{}
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
_, err = as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.Error(t, err)
//
cfg = &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "by-dev", TargetClusterId: "test2"},
},
}
// Test update pass.
_, err = as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.NoError(t, err)
// Test idempotent
b.EXPECT().GetLatestChannelAssignment().Unset()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: cfg,
}, nil).Maybe()
_, err = as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.NoError(t, err)
// Test secondary path.
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).Unset()
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, rk ...message.ResourceKey) (broadcaster.BroadcastAPI, error) {
return nil, broadcaster.ErrNotPrimary
})
// Still idempotent.
_, err = as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.NoError(t, err)
// Test update on secondary path, it should be block until the replicate configuration is changed.
b.EXPECT().GetLatestChannelAssignment().Unset()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "test2", TargetClusterId: "by-dev"},
},
},
}, nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).Unset()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
select {
case <-time.After(500 * time.Millisecond):
return cb(balancer.WatchChannelAssignmentsCallbackParam{
ReplicateConfiguration: cfg,
})
case <-ctx.Done():
return ctx.Err()
}
})
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
_, err = as.UpdateReplicateConfiguration(ctx, &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.ErrorIs(t, err, context.DeadlineExceeded)
_, err = as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.NoError(t, err)
// Test callback
b.EXPECT().UpdateReplicateConfiguration(mock.Anything, mock.Anything).Return(nil)
msg := message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: cfg,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithBroadcast([]string{"v1"}).
MustBuildBroadcast()
assert.NoError(t, registry.CallMessageAckCallback(context.Background(), msg, map[string]*message.AppendResult{}))
}
func TestForcePromoteValidation(t *testing.T) {
t.Run("valid_single_cluster_no_topology", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.NoError(t, err)
})
// Test validateForcePromoteConfiguration standalone function directly
t.Run("valid_single_cluster_no_topology_direct", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.NoError(t, err)
})
t.Run("invalid_multiple_clusters", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{},
}
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.Error(t, err)
// Config helper validates this as "primary count is not 1"
assert.Contains(t, err.Error(), "primary count is not 1")
})
t.Run("invalid_with_topology_direct", func(t *testing.T) {
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "by-dev", TargetClusterId: "test2"},
},
}
err := validateForcePromoteConfiguration(cfg, "by-dev")
assert.Error(t, err)
})
}
func TestForcePromoteOnPrimaryCluster(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Current cluster is a primary
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
balance.Register(b)
// Set up broadcaster to return ErrNotSecondary (we are primary, force promote should fail)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(nil, broadcaster.ErrNotSecondary).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote on a primary cluster should fail - use empty config
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{},
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "force promote can only be used on secondary clusters")
}
func TestForcePromoteSuccess(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Current cluster is secondary (has topology from primary to by-dev)
currentReplicateConfig := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "primary", Pchannels: []string{"primary-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://primary:19530", Token: "primary"}},
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "primary", TargetClusterId: "by-dev"},
},
}
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: currentReplicateConfig,
}, nil).Maybe()
balance.Register(b)
// Set up broadcaster with WithSecondaryClusterResourceKey (we are secondary)
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(&types.BroadcastAppendResult{
BroadcastID: 1,
AppendResults: map[string]*types.AppendResult{
"by-dev-1": {TimeTick: 100},
},
}, nil).Maybe()
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with config specifying only the current cluster (no topology)
resp, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.NoError(t, err)
assert.NotNil(t, resp)
}
func TestForcePromoteIdempotent(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Config already applied - standalone primary with matching pchannel
// This will be matched against the auto-constructed config
alreadyPromotedConfig := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{
ClusterId: "by-dev",
Pchannels: []string{"by-dev-1"},
ConnectionParam: &commonpb.ConnectionParam{
Uri: "http://localhost:19530",
Token: "",
},
},
},
}
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: alreadyPromotedConfig,
}, nil).Maybe()
balance.Register(b)
// Set up broadcaster with WithSecondaryClusterResourceKey
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with config specifying only the current cluster (no topology)
resp, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.NoError(t, err)
assert.NotNil(t, resp)
}
func TestAlterReplicateConfigCallbackForcePromote(t *testing.T) {
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
// Set up balancer
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().UpdateReplicateConfiguration(mock.Anything, mock.Anything).Return(nil).Maybe()
balance.Register(b)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
as := NewAssignmentService()
_ = as
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
}
// Test callback with force promote flag set
msg := message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: cfg,
ForcePromote: true,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithBroadcast([]string{"v1"}).
MustBuildBroadcast()
assert.NoError(t, registry.CallMessageAckCallback(context.Background(), msg, map[string]*message.AppendResult{}))
// Test callback without force promote
msg2 := message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: cfg,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithBroadcast([]string{"v1"}).
MustBuildBroadcast()
assert.NoError(t, registry.CallMessageAckCallback(context.Background(), msg2, map[string]*message.AppendResult{}))
}
func TestAlterReplicateConfigCallbackIgnore(t *testing.T) {
// Test callback with ignore flag set - should skip all processing
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
// Set up balancer - should NOT be called
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Note: UpdateReplicateConfiguration should NOT be called because ignore=true
balance.Register(b)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
as := NewAssignmentService()
_ = as
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
}
// Test callback with ignore flag set - should skip all processing
msg := message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: cfg,
ForcePromote: true,
Ignore: true,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithBroadcast([]string{"v1"}).
MustBuildBroadcast()
assert.NoError(t, registry.CallMessageAckCallback(context.Background(), msg, map[string]*message.AppendResult{}))
}
func TestForcePromoteBroadcastOtherError(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
balance.Register(b)
// Broadcaster returns a non-ErrNotSecondary error
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(nil, errors.New("broadcaster unavailable"))
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with empty config
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{},
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "broadcaster unavailable")
}
func TestForcePromoteBroadcastAppendError(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
currentReplicateConfig := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "primary", Pchannels: []string{"primary-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://primary:19530", Token: "primary"}},
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "primary", TargetClusterId: "by-dev"},
},
}
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: currentReplicateConfig,
}, nil).Maybe()
balance.Register(b)
// Set up broadcaster with WithSecondaryClusterResourceKey returning broadcast API that fails on Broadcast
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(nil, errors.New("broadcast append failed"))
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with config specifying only the current cluster (no topology)
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "broadcast append failed")
}
func TestAlterReplicateConfigCallbackErrors(t *testing.T) {
resource.InitForTest()
t.Run("update_config_error", func(t *testing.T) {
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().UpdateReplicateConfiguration(mock.Anything, mock.Anything).Return(errors.New("update failed"))
balance.Register(b)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
as := NewAssignmentService()
_ = as
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
}
msg := message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: cfg,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithBroadcast([]string{"v1"}).
MustBuildBroadcast()
err := registry.CallMessageAckCallback(context.Background(), msg, map[string]*message.AppendResult{})
assert.Error(t, err)
assert.Contains(t, err.Error(), "update failed")
})
}
func TestUpdateReplicateConfigNonPrimaryBroadcastError(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
// Mock channel.GetClusterChannels to avoid blocking on unregistered singleton.
mockGetClusterChannels := mockey.Mock(channel.GetClusterChannels).Return(message.ClusterChannels{
Channels: []string{"by-dev-1"},
ControlChannel: "by-dev-1_vcchan",
}).Build()
defer mockGetClusterChannels.UnPatch()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
balance.Register(b)
// Broadcaster returns a non-ErrNotPrimary error (not force promote path)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).Return(nil, errors.New("unexpected error"))
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "by-dev", TargetClusterId: "test2"},
},
}
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "unexpected error")
}
func TestUpdateReplicateConfigBroadcastError(t *testing.T) {
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
// Mock channel.GetClusterChannels to avoid blocking on unregistered singleton.
mockGetClusterChannels := mockey.Mock(channel.GetClusterChannels).Return(message.ClusterChannels{
Channels: []string{"by-dev-1"},
ControlChannel: "by-dev-1_vcchan",
}).Build()
defer mockGetClusterChannels.UnPatch()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil).Maybe()
balance.Register(b)
// Broadcaster acquires lock but Broadcast() fails
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(nil, errors.New("broadcast failed"))
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).Return(mba, nil)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "by-dev", TargetClusterId: "test2"},
},
}
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "broadcast failed")
}
func TestUpdateReplicateConfigSecondValidateSameConfig(t *testing.T) {
// Tests the path where the second validateReplicateConfiguration returns errReplicateConfigurationSame
// (config became same between first check and getting the lock)
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
// Mock channel.GetClusterChannels to avoid blocking on unregistered singleton.
mockGetClusterChannels := mockey.Mock(channel.GetClusterChannels).Return(message.ClusterChannels{
Channels: []string{"by-dev-1"},
ControlChannel: "by-dev-1_vcchan",
}).Build()
defer mockGetClusterChannels.UnPatch()
callCount := 0
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "by-dev", TargetClusterId: "test2"},
},
}
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().RunAndReturn(func() (*balancer.WatchChannelAssignmentsCallbackParam, error) {
callCount++
if callCount <= 1 {
// First call: config is different (nil)
return &balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil
}
// Second call: config is now the same (was applied by another path)
return &balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: cfg,
}, nil
})
balance.Register(b)
// Broadcaster acquires lock successfully
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).Return(mba, nil)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
resp, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.NoError(t, err)
assert.NotNil(t, resp)
}
func TestAssignmentDiscoverBalanceError(t *testing.T) {
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager() // Reset so GetWithContext fails with canceled context
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
ctx, cancel := context.WithCancel(context.Background())
cancel() // cancel immediately
ss := mock_streamingpb.NewMockStreamingCoordAssignmentService_AssignmentDiscoverServer(t)
ss.EXPECT().Context().Return(ctx).Maybe()
err := as.AssignmentDiscover(ss)
assert.Error(t, err)
}
func TestUpdateWALBalancePolicyError(t *testing.T) {
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager() // Reset so GetWithContext fails with canceled context
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
ctx, cancel := context.WithCancel(context.Background())
cancel() // cancel immediately so GetWithContext fails
_, err := as.UpdateWALBalancePolicy(ctx, &streamingpb.UpdateWALBalancePolicyRequest{})
assert.Error(t, err)
}
func TestWaitUntilPrimaryChangeBalanceError(t *testing.T) {
// Covers waitUntilPrimaryChangeOrConfigurationSame balance.GetWithContext error (line 123-125)
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager() // Reset so GetWithContext fails
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := &assignmentServiceImpl{}
ctx, cancel := context.WithCancel(context.Background())
cancel()
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
}
err := as.waitUntilPrimaryChangeOrConfigurationSame(ctx, cfg)
assert.Error(t, err)
}
func TestHandleForcePromoteBalanceError(t *testing.T) {
// Covers handleForcePromote balance.GetWithContext error
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager() // Reset so GetWithContext fails
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
ctx, cancel := context.WithCancel(context.Background())
cancel()
// Force promote with empty config
_, err := as.UpdateReplicateConfiguration(ctx, &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{},
ForcePromote: true,
})
assert.Error(t, err)
}
func TestHandleForcePromoteGetAssignmentError(t *testing.T) {
// Covers handleForcePromote GetLatestChannelAssignment error
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// GetLatestChannelAssignment in handleForcePromote fails
b.EXPECT().GetLatestChannelAssignment().Return(nil, errors.New("assignment error")).Maybe()
balance.Register(b)
// WithSecondaryClusterResourceKey succeeds (we are secondary)
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with config specifying only the current cluster (no topology)
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "assignment error")
}
func TestHandleForcePromoteSameConfigAfterBroadcasterCheck(t *testing.T) {
// Covers handleForcePromote double-check same config - auto-constructed config matches existing
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
// The auto-constructed standalone primary config
// ConnectionParam must match what exists in the secondary config
forcePromoteCfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{
ClusterId: "by-dev",
Pchannels: []string{"by-dev-1"},
ConnectionParam: &commonpb.ConnectionParam{
Uri: "http://test:19530",
Token: "by-dev",
},
},
},
}
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Return the standalone primary config directly
// This tests the idempotent path where config already matches
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: forcePromoteCfg,
}, nil).Maybe()
balance.Register(b)
// Broadcaster returns success for secondary check
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with config specifying only the current cluster (no topology)
resp, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.NoError(t, err)
assert.NotNil(t, resp)
}
func TestHandleForcePromoteValidatorError(t *testing.T) {
// Covers handleForcePromote validator.Validate() error
// The validator may fail if pchannels in cluster view don't match cluster config
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Return config with extra pchannel in PChannelView that causes validator to fail
// PChannelView has 2 pchannels but current config's cluster has only 1
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
{Name: "other-1"}: channel.NewPChannelMeta("other-1", types.AccessModeRW),
},
},
ReplicateConfiguration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "primary", Pchannels: []string{"primary-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://primary:19530", Token: "primary"}},
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "primary", TargetClusterId: "by-dev"},
},
},
}, nil).Maybe()
balance.Register(b)
// WithSecondaryClusterResourceKey succeeds
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with config specifying only the current cluster (no topology)
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.Error(t, err)
}
func TestHandleForcePromoteNoCurrentConfig(t *testing.T) {
// Covers handleForcePromote when current config is nil
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Return assignment with nil ReplicateConfiguration
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: nil, // No config exists
}, nil).Maybe()
balance.Register(b)
// WithSecondaryClusterResourceKey succeeds
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "force promote requires existing replicate configuration")
}
func TestHandleForcePromoteClusterNotFound(t *testing.T) {
// Covers handleForcePromote when current cluster is not found in config
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
// Return config with different cluster ID (not "by-dev")
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
ReplicateConfiguration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "other-cluster", Pchannels: []string{"other-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://other:19530", Token: "other"}},
},
},
}, nil).Maybe()
balance.Register(b)
// WithSecondaryClusterResourceKey succeeds
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "force promote requires current cluster in existing configuration")
}
func TestSecondValidateNonSameError(t *testing.T) {
// Covers line 111: second validateReplicateConfiguration returns non-same error
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
// Mock channel.GetClusterChannels to avoid blocking on unregistered singleton.
mockGetClusterChannels := mockey.Mock(channel.GetClusterChannels).Return(message.ClusterChannels{
Channels: []string{"by-dev-1"},
ControlChannel: "by-dev-1_vcchan",
}).Build()
defer mockGetClusterChannels.UnPatch()
callCount := 0
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
b.EXPECT().GetLatestChannelAssignment().RunAndReturn(func() (*balancer.WatchChannelAssignmentsCallbackParam, error) {
callCount++
if callCount <= 1 {
return &balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
},
},
}, nil
}
// Second call after lock: return error
return nil, errors.New("assignment unavailable")
})
balance.Register(b)
// Broadcaster lock succeeds
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithResourceKeys(mock.Anything, mock.Anything).Return(mba, nil)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
{ClusterId: "test2", Pchannels: []string{"test2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test2:19530", Token: "test2"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "by-dev", TargetClusterId: "test2"},
},
}
_, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: cfg,
})
assert.Error(t, err)
assert.Contains(t, err.Error(), "assignment unavailable")
}
func TestAlterReplicateConfigCallbackBalanceError(t *testing.T) {
// Covers alterReplicateConfiguration balance.GetWithContext error (lines 345-347)
resource.InitForTest()
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
_ = as
cfg := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
}
msg := message.NewAlterReplicateConfigMessageBuilderV2().
WithHeader(&message.AlterReplicateConfigMessageHeader{
ReplicateConfiguration: cfg,
}).
WithBody(&message.AlterReplicateConfigMessageBody{}).
WithBroadcast([]string{"v1"}).
MustBuildBroadcast()
ctx, cancel := context.WithCancel(context.Background())
cancel()
err := registry.CallMessageAckCallback(ctx, msg, map[string]*message.AppendResult{})
assert.Error(t, err)
}
func TestForcePromoteMultiplePChannels(t *testing.T) {
// Tests force promote with multiple pchannels - auto-constructed config includes all pchannels
resource.InitForTest()
mw := mock_streaming.NewMockWALAccesser(t)
mw.EXPECT().ControlChannel().Return("by-dev-1_vcchan").Maybe()
streaming.SetWALForTest(mw)
broadcast.ResetBroadcaster()
snmanager.ResetStreamingNodeManager()
b := mock_balancer.NewMockBalancer(t)
b.EXPECT().WaitUntilWALbasedDDLReady(mock.Anything).Return(nil).Maybe()
b.EXPECT().WatchChannelAssignments(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, cb balancer.WatchChannelAssignmentsCallback) error {
<-ctx.Done()
return ctx.Err()
}).Maybe()
b.EXPECT().Close().Return().Maybe()
currentReplicateConfig := &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "primary", Pchannels: []string{"primary-1", "primary-2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://primary:19530", Token: "primary"}},
{ClusterId: "by-dev", Pchannels: []string{"by-dev-1", "by-dev-2"}, ConnectionParam: &commonpb.ConnectionParam{Uri: "http://test:19530", Token: "by-dev"}},
},
CrossClusterTopology: []*commonpb.CrossClusterTopology{
{SourceClusterId: "primary", TargetClusterId: "by-dev"},
},
}
// Multiple pchannels - by-dev-1 matches control channel, by-dev-2 does not
b.EXPECT().GetLatestChannelAssignment().Return(&balancer.WatchChannelAssignmentsCallbackParam{
PChannelView: &channel.PChannelView{
Channels: map[channel.ChannelID]*channel.PChannelMeta{
{Name: "by-dev-1"}: channel.NewPChannelMeta("by-dev-1", types.AccessModeRW),
{Name: "by-dev-2"}: channel.NewPChannelMeta("by-dev-2", types.AccessModeRW),
},
},
ReplicateConfiguration: currentReplicateConfig,
}, nil).Maybe()
balance.Register(b)
// Set up broadcaster with WithSecondaryClusterResourceKey
mba := mock_broadcaster.NewMockBroadcastAPI(t)
mba.EXPECT().Broadcast(mock.Anything, mock.Anything).Return(&types.BroadcastAppendResult{
BroadcastID: 1,
AppendResults: map[string]*types.AppendResult{
"by-dev-1": {TimeTick: 100},
"by-dev-2": {TimeTick: 101},
},
}, nil).Maybe()
mba.EXPECT().Close().Return().Maybe()
mb := mock_broadcaster.NewMockBroadcaster(t)
mb.EXPECT().WithSecondaryClusterResourceKey(mock.Anything).Return(mba, nil).Maybe()
mb.EXPECT().Close().Return().Maybe()
broadcast.Register(mb)
as := NewAssignmentService()
// Force promote with config specifying only the current cluster (no topology)
resp, err := as.UpdateReplicateConfiguration(context.Background(), &streamingpb.UpdateReplicateConfigurationRequest{
Configuration: &commonpb.ReplicateConfiguration{
Clusters: []*commonpb.MilvusCluster{
{ClusterId: "by-dev"},
},
},
ForcePromote: true,
})
assert.NoError(t, err)
assert.NotNil(t, resp)
}