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)
}