package action

import (
	"reflect"
	"testing"

	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
	apiv1 "github.com/tidbcloud/cluster-service-ng/api/cluster-service-ng/v1beta"
	dplanemgrv1 "github.com/tidbcloud/cluster-service-ng/api/dataplane-manager/v1"
	"github.com/tidbcloud/cluster-service-ng/domain/entity"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow"
	dataplanemgrmocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	mysqlschema "github.com/tidbcloud/cluster-service-ng/pkg/storage/mysql/schema"
	repomocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/encoding/protojson"
	"google.golang.org/protobuf/types/known/emptypb"
)

func TestBuildPublicLoadBalancer(t *testing.T) {
	tests := []struct {
		name                     string
		gatewayIdentity          string
		disablePublicDNSIdentity bool
		sharedTiProxy            bool
		expectedIdentity         string
		expectedSubIdentity      string
		expectPreserveClientIP   *bool
	}{
		{
			name:                "uses gateway identity and dns identity by default",
			gatewayIdentity:     "cluster-gateway",
			expectedIdentity:    "cluster-gateway",
			expectedSubIdentity: "group-dns",
		},
		{
			name:                     "drops public dns sub identity when disabled",
			gatewayIdentity:          "cluster-gateway",
			disablePublicDNSIdentity: true,
			expectedIdentity:         "cluster-gateway",
			expectedSubIdentity:      "",
		},
		{
			name:                     "keeps old fallback when gateway identity is empty",
			disablePublicDNSIdentity: true,
			expectedIdentity:         "group-dns",
			expectedSubIdentity:      "",
		},
		{
			name:                   "shared tiproxy public lb disables preserve client ip",
			gatewayIdentity:        "cluster-gateway",
			sharedTiProxy:          true,
			expectedIdentity:       "cluster-gateway",
			expectedSubIdentity:    "group-dns",
			expectPreserveClientIP: new(bool),
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tidbCluster := &entity.TiDBCluster{GatewayIdentity: tt.gatewayIdentity}
			tidbGroup := &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id:           "tidb-group-1",
					SharedPoolId: "shared-pool-1",
					SqlEndpoint: &entity.SQLEndpoint{
						DnsIdentity:              "group-dns",
						DisablePublicDnsIdentity: tt.disablePublicDNSIdentity,
						PublicLb: &entity.LoadBalancer{
							Enabled: true,
							Port:    4000,
						},
					},
				},
			}
			lbComp := &entity.Component{
				Id:            "tiproxy-1",
				ComponentType: apiv1.ComponentType_TIPROXY,
			}
			sharedPool := &entity.SharedPool{}
			if tt.sharedTiProxy {
				sharedPool.DeployType = mysqlschema.SharedPoolDeployTypeSharedTiProxy
			}

			lb := buildPublicLoadBalancer(tidbCluster, tidbGroup, lbComp, sharedPool)

			require.Equal(t, tt.expectedIdentity, lb.GetIdentity())
			require.Equal(t, tt.expectedSubIdentity, lb.GetSubIdentity())
			require.Equal(t, tt.expectPreserveClientIP, lb.PreserveClientIp)
		})
	}
}

func TestBuildLoadBalancersGracefulRestartAnnotations(t *testing.T) {
	tidbCluster := &entity.TiDBCluster{GatewayIdentity: "cluster-gateway"}
	lbComp := &entity.Component{
		Id:            "tiproxy-1",
		ComponentType: apiv1.ComponentType_TIPROXY,
	}

	tests := []struct {
		name              string
		provider          apiv1.ProviderType
		publicPort        int32
		privatePort       int32
		expectAnnotations bool
	}{
		{
			name:              "aws",
			provider:          apiv1.ProviderType_AWS,
			publicPort:        4000,
			privatePort:       4000,
			expectAnnotations: true,
		},
		{
			name:              "alicloud",
			provider:          apiv1.ProviderType_ALICLOUD,
			publicPort:        4000,
			privatePort:       4000,
			expectAnnotations: true,
		},
		{
			name:              "azure uses per-port annotations",
			provider:          apiv1.ProviderType_AZURE,
			publicPort:        4001,
			privatePort:       4000,
			expectAnnotations: true,
		},
		{
			name:              "unsupported provider",
			provider:          apiv1.ProviderType_GCP,
			publicPort:        4000,
			privatePort:       4000,
			expectAnnotations: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			tidbGroup := &entity.TiDBGroup{
				TiDBGroup: apiv1.TiDBGroup{
					Id:           "tidb-group-1",
					SharedPoolId: "shared-pool-1",
					SqlEndpoint: &entity.SQLEndpoint{
						DnsIdentity: "group-dns",
						PublicLb:    &entity.LoadBalancer{Enabled: true, Port: tt.publicPort},
						PrivateLb:   &entity.LoadBalancer{Enabled: true, Port: tt.privatePort},
					},
				},
			}
			sharedPool := &entity.SharedPool{
				SharedPool: apiv1.SharedPool{Provider: &apiv1.Provider{Type: tt.provider}},
			}

			publicLB := buildPublicLoadBalancer(tidbCluster, tidbGroup, lbComp, sharedPool)
			privateLB := buildPrivateLoadBalancer(tidbCluster, tidbGroup, lbComp, sharedPool)

			if tt.expectAnnotations {
				expectedPublic := entity.BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool, tt.publicPort)
				expectedPrivate := entity.BuildTiProxyGracefulRestartLoadBalancerAnnotations(sharedPool, tt.privatePort)
				require.Equal(t, expectedPublic, publicLB.GetAnnotations())
				require.Equal(t, expectedPrivate, privateLB.GetAnnotations())
				require.True(t, publicLB.GetGracefulUnhealthyEnabled())
				require.True(t, privateLB.GetGracefulUnhealthyEnabled())
			} else {
				require.Empty(t, publicLB.GetAnnotations())
				require.Empty(t, privateLB.GetAnnotations())
				require.False(t, publicLB.GetGracefulUnhealthyEnabled())
				require.False(t, privateLB.GetGracefulUnhealthyEnabled())
			}
		})
	}
}

func TestBuildPrivateLoadBalancerUsesConfiguredIpAllowList(t *testing.T) {
	tidbCluster := &entity.TiDBCluster{GatewayIdentity: "cluster-gateway"}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:           "tidb-group-1",
			SharedPoolId: "shared-pool-1",
			SqlEndpoint: &entity.SQLEndpoint{
				DnsIdentity: "group-dns",
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "192.168.0.0/24"},
						{CidrBlock: "10.0.0.0/16"},
					},
				},
			},
		},
	}
	lbComp := &entity.Component{
		Id:            "tiproxy-1",
		ComponentType: apiv1.ComponentType_TIPROXY,
	}

	lb := buildPrivateLoadBalancer(tidbCluster, tidbGroup, lbComp, &entity.SharedPool{})

	require.Equal(t, []string{"192.168.0.0/24", "10.0.0.0/16"}, lb.GetIpWhiteList())
	require.Nil(t, lb.PreserveClientIp)
}

func TestBuildPrivateLoadBalancerSkipsNilAndEmptyIpAllowListEntries(t *testing.T) {
	tidbCluster := &entity.TiDBCluster{GatewayIdentity: "cluster-gateway"}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:           "tidb-group-1",
			SharedPoolId: "shared-pool-1",
			SqlEndpoint: &entity.SQLEndpoint{
				DnsIdentity: "group-dns",
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						nil,
						{CidrBlock: ""},
						{CidrBlock: "192.168.0.0/24"},
					},
				},
			},
		},
	}
	lbComp := &entity.Component{
		Id:            "tiproxy-1",
		ComponentType: apiv1.ComponentType_TIPROXY,
	}

	lb := buildPrivateLoadBalancer(tidbCluster, tidbGroup, lbComp, &entity.SharedPool{})

	require.Equal(t, []string{"192.168.0.0/24"}, lb.GetIpWhiteList())
	require.Nil(t, lb.PreserveClientIp)
}

func TestApplyLoadBalanceActionUpdatesExistingPrivateLoadBalancer(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	tidbGroupRepo := repomocks.NewTiDBGroupRepo(t)
	tidbClusterRepo := repomocks.NewTiDBClusterRepo(t)
	client := dataplanemgrmocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "shared-pool-1",
			DataplaneName: "dataplane-1",
			ComponentGroups: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
		},
	}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:            "tidb-group-1",
			ClusterId:     "cluster-1",
			SharedPoolId:  "shared-pool-1",
			DataplaneName: "dataplane-1",
			SqlEndpoint: &entity.SQLEndpoint{
				PublicLb: &entity.LoadBalancer{Enabled: false},
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "192.168.0.0/24"},
						{CidrBlock: "10.0.0.0/16"},
					},
				},
			},
		},
	}
	tidbCluster := &entity.TiDBCluster{
		Id:              "cluster-1",
		GatewayIdentity: "cluster-gateway",
	}

	tidbGroupStr, err := protojson.Marshal(tidbGroup)
	require.NoError(t, err)

	sharedPoolRepo.On("Load", mock.Anything, "shared-pool-1", false).Return(sharedPool, nil).Once()
	tidbGroupRepo.On("Load", mock.Anything, "tidb-group-1", false).Return(tidbGroup, nil).Once()
	tidbClusterRepo.On("Load", mock.Anything, "cluster-1", false).Return(tidbCluster, nil).Once()
	client.On("GetDataPlane", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetDataPlaneRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1"
	})).Return(&dplanemgrv1.DataPlane{
		Vpc: &dplanemgrv1.DataPlane_Vpc{Cidr: "10.0.0.0/16"},
	}, nil).Once()

	client.On("DeleteLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.DeleteLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" && req.GetId() == entity.GetPublicSvcName("tidb-group-1")
	})).Return(&emptypb.Empty{}, nil).Once()
	client.On("GetLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" && req.GetId() == entity.GetPrivateSvcName("tidb-group-1")
	})).Return(&dplanemgrv1.LoadBalancer{
		Id:                       entity.GetPrivateSvcName("tidb-group-1"),
		GracefulUnhealthyEnabled: true,
	}, nil).Once()
	client.On("UpdateLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.UpdateLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" &&
			req.GetId() == entity.GetPrivateSvcName("tidb-group-1") &&
			req.GetLoadbalancer().GetPort() == 4000 &&
			req.GetLoadbalancer().GetType() == dplanemgrv1.LoadBalancer_PRIVATE &&
			!req.GetLoadbalancer().GetGracefulUnhealthyEnabled() &&
			len(req.GetLoadbalancer().GetAnnotations()) == 0 &&
			reflect.DeepEqual(req.GetLoadbalancer().GetIpWhiteList(), []string{"192.168.0.0/24", "10.0.0.0/16"})
	})).Return(&dplanemgrv1.LoadBalancer{}, nil).Once()

	action := &ApplyLoadBalanceAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbGroupRepo:   tidbGroupRepo,
		tidbClusterRepo: tidbClusterRepo,
		client:          client,
		dataplaneClient: client,
	}

	err = action.Invoke(actions.ActionArgs{
		workflow.UpdateLoadBalancerKey: actions.ActionArgString(workflow.UpdateLoadBalancerValue),
		workflow.SharedPoolIdKey:       actions.ActionArgString("shared-pool-1"),
		workflow.TiDBGroupIdKey:        actions.ActionArgString("tidb-group-1"),
		workflow.TiDBGroupKey:          actions.ActionArgString(string(tidbGroupStr)),
	})
	require.NoError(t, err)
}

func TestApplyLoadBalanceActionCreatesPrivateLoadBalancerWhenMissing(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	tidbGroupRepo := repomocks.NewTiDBGroupRepo(t)
	tidbClusterRepo := repomocks.NewTiDBClusterRepo(t)
	client := dataplanemgrmocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "shared-pool-1",
			DataplaneName: "dataplane-1",
			ComponentGroups: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
		},
	}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:            "tidb-group-1",
			ClusterId:     "cluster-1",
			SharedPoolId:  "shared-pool-1",
			DataplaneName: "dataplane-1",
			SqlEndpoint: &entity.SQLEndpoint{
				PublicLb: &entity.LoadBalancer{Enabled: false},
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
					IpAllowList: []*entity.IpAllowList{
						{CidrBlock: "192.168.0.0/24"},
					},
				},
			},
		},
	}
	tidbCluster := &entity.TiDBCluster{
		Id:              "cluster-1",
		GatewayIdentity: "cluster-gateway",
	}

	tidbGroupStr, err := protojson.Marshal(tidbGroup)
	require.NoError(t, err)

	sharedPoolRepo.On("Load", mock.Anything, "shared-pool-1", false).Return(sharedPool, nil).Once()
	tidbGroupRepo.On("Load", mock.Anything, "tidb-group-1", false).Return(tidbGroup, nil).Once()
	tidbClusterRepo.On("Load", mock.Anything, "cluster-1", false).Return(tidbCluster, nil).Once()
	client.On("GetDataPlane", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetDataPlaneRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1"
	})).Return(&dplanemgrv1.DataPlane{
		Vpc: &dplanemgrv1.DataPlane_Vpc{Cidr: "10.0.0.0/16"},
	}, nil).Once()

	client.On("DeleteLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.DeleteLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" && req.GetId() == entity.GetPublicSvcName("tidb-group-1")
	})).Return(&emptypb.Empty{}, nil).Once()
	client.On("GetLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" && req.GetId() == entity.GetPrivateSvcName("tidb-group-1")
	})).Return(nil, status.Error(codes.NotFound, "missing")).Once()
	client.On("CreateLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.CreateLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" &&
			req.GetId() == entity.GetPrivateSvcName("tidb-group-1") &&
			reflect.DeepEqual(req.GetLoadbalancer().GetIpWhiteList(), []string{"192.168.0.0/24", "10.0.0.0/16"})
	})).Return(&dplanemgrv1.LoadBalancer{}, nil).Once()

	action := &ApplyLoadBalanceAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbGroupRepo:   tidbGroupRepo,
		tidbClusterRepo: tidbClusterRepo,
		client:          client,
		dataplaneClient: client,
	}

	err = action.Invoke(actions.ActionArgs{
		workflow.UpdateLoadBalancerKey: actions.ActionArgString(workflow.UpdateLoadBalancerValue),
		workflow.SharedPoolIdKey:       actions.ActionArgString("shared-pool-1"),
		workflow.TiDBGroupIdKey:        actions.ActionArgString("tidb-group-1"),
		workflow.TiDBGroupKey:          actions.ActionArgString(string(tidbGroupStr)),
	})
	require.NoError(t, err)
}

func TestApplyLoadBalanceActionUpdatesExistingPrivateLoadBalancerWhenMetadataAllowlistEmpty(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	tidbGroupRepo := repomocks.NewTiDBGroupRepo(t)
	tidbClusterRepo := repomocks.NewTiDBClusterRepo(t)
	client := dataplanemgrmocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "shared-pool-1",
			DataplaneName: "dataplane-1",
			ComponentGroups: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
		},
	}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:            "tidb-group-1",
			ClusterId:     "cluster-1",
			SharedPoolId:  "shared-pool-1",
			DataplaneName: "dataplane-1",
			SqlEndpoint: &entity.SQLEndpoint{
				PublicLb: &entity.LoadBalancer{Enabled: false},
				PrivateLb: &entity.LoadBalancer{
					Enabled:     true,
					Port:        4000,
					IpAllowList: nil,
				},
			},
		},
	}
	tidbCluster := &entity.TiDBCluster{
		Id:              "cluster-1",
		GatewayIdentity: "cluster-gateway",
	}

	tidbGroupStr, err := protojson.Marshal(tidbGroup)
	require.NoError(t, err)

	sharedPoolRepo.On("Load", mock.Anything, "shared-pool-1", false).Return(sharedPool, nil).Once()
	tidbGroupRepo.On("Load", mock.Anything, "tidb-group-1", false).Return(tidbGroup, nil).Once()
	tidbClusterRepo.On("Load", mock.Anything, "cluster-1", false).Return(tidbCluster, nil).Once()
	client.On("GetDataPlane", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetDataPlaneRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1"
	})).Return(&dplanemgrv1.DataPlane{
		Vpc: &dplanemgrv1.DataPlane_Vpc{Cidr: "10.0.0.0/16"},
	}, nil).Once()

	client.On("DeleteLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.DeleteLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" && req.GetId() == entity.GetPublicSvcName("tidb-group-1")
	})).Return(&emptypb.Empty{}, nil).Once()
	client.On("GetLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" && req.GetId() == entity.GetPrivateSvcName("tidb-group-1")
	})).Return(&dplanemgrv1.LoadBalancer{Id: entity.GetPrivateSvcName("tidb-group-1")}, nil).Once()
	client.On("UpdateLoadBalancer", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.UpdateLoadBalancerRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1" &&
			req.GetId() == entity.GetPrivateSvcName("tidb-group-1") &&
			req.GetLoadbalancer().GetPort() == 4000 &&
			req.GetLoadbalancer().GetType() == dplanemgrv1.LoadBalancer_PRIVATE &&
			reflect.DeepEqual(req.GetLoadbalancer().GetIpWhiteList(), []string{"10.0.0.0/16"})
	})).Return(&dplanemgrv1.LoadBalancer{}, nil).Once()

	action := &ApplyLoadBalanceAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbGroupRepo:   tidbGroupRepo,
		tidbClusterRepo: tidbClusterRepo,
		client:          client,
		dataplaneClient: client,
	}

	err = action.Invoke(actions.ActionArgs{
		workflow.UpdateLoadBalancerKey: actions.ActionArgString(workflow.UpdateLoadBalancerValue),
		workflow.SharedPoolIdKey:       actions.ActionArgString("shared-pool-1"),
		workflow.TiDBGroupIdKey:        actions.ActionArgString("tidb-group-1"),
		workflow.TiDBGroupKey:          actions.ActionArgString(string(tidbGroupStr)),
	})
	require.NoError(t, err)
}

func TestApplyLoadBalanceActionReturnsErrorWhenGetDataPlaneFails(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	tidbGroupRepo := repomocks.NewTiDBGroupRepo(t)
	tidbClusterRepo := repomocks.NewTiDBClusterRepo(t)
	client := dataplanemgrmocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "shared-pool-1",
			DataplaneName: "dataplane-1",
			ComponentGroups: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
		},
	}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:            "tidb-group-1",
			ClusterId:     "cluster-1",
			SharedPoolId:  "shared-pool-1",
			DataplaneName: "dataplane-1",
			SqlEndpoint: &entity.SQLEndpoint{
				PublicLb: &entity.LoadBalancer{Enabled: false},
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
				},
			},
		},
	}
	tidbCluster := &entity.TiDBCluster{
		Id:              "cluster-1",
		GatewayIdentity: "cluster-gateway",
	}

	tidbGroupStr, err := protojson.Marshal(tidbGroup)
	require.NoError(t, err)

	sharedPoolRepo.On("Load", mock.Anything, "shared-pool-1", false).Return(sharedPool, nil).Once()
	tidbGroupRepo.On("Load", mock.Anything, "tidb-group-1", false).Return(tidbGroup, nil).Once()
	tidbClusterRepo.On("Load", mock.Anything, "cluster-1", false).Return(tidbCluster, nil).Once()
	client.On("GetDataPlane", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetDataPlaneRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1"
	})).Return(nil, status.Error(codes.Internal, "dataplane lookup failed")).Once()

	action := &ApplyLoadBalanceAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbGroupRepo:   tidbGroupRepo,
		tidbClusterRepo: tidbClusterRepo,
		client:          client,
		dataplaneClient: client,
	}

	err = action.Invoke(actions.ActionArgs{
		workflow.UpdateLoadBalancerKey: actions.ActionArgString(workflow.UpdateLoadBalancerValue),
		workflow.SharedPoolIdKey:       actions.ActionArgString("shared-pool-1"),
		workflow.TiDBGroupIdKey:        actions.ActionArgString("tidb-group-1"),
		workflow.TiDBGroupKey:          actions.ActionArgString(string(tidbGroupStr)),
	})
	require.ErrorContains(t, err, "failed to get dataplane dataplane-1")
}

func TestApplyLoadBalanceActionReturnsErrorWhenDataplaneCIDREmpty(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	tidbGroupRepo := repomocks.NewTiDBGroupRepo(t)
	tidbClusterRepo := repomocks.NewTiDBClusterRepo(t)
	client := dataplanemgrmocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "shared-pool-1",
			DataplaneName: "dataplane-1",
			ComponentGroups: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
		},
	}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:            "tidb-group-1",
			ClusterId:     "cluster-1",
			SharedPoolId:  "shared-pool-1",
			DataplaneName: "dataplane-1",
			SqlEndpoint: &entity.SQLEndpoint{
				PublicLb: &entity.LoadBalancer{Enabled: false},
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
				},
			},
		},
	}
	tidbCluster := &entity.TiDBCluster{
		Id:              "cluster-1",
		GatewayIdentity: "cluster-gateway",
	}

	tidbGroupStr, err := protojson.Marshal(tidbGroup)
	require.NoError(t, err)

	sharedPoolRepo.On("Load", mock.Anything, "shared-pool-1", false).Return(sharedPool, nil).Once()
	tidbGroupRepo.On("Load", mock.Anything, "tidb-group-1", false).Return(tidbGroup, nil).Once()
	tidbClusterRepo.On("Load", mock.Anything, "cluster-1", false).Return(tidbCluster, nil).Once()
	client.On("GetDataPlane", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetDataPlaneRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1"
	})).Return(&dplanemgrv1.DataPlane{
		Vpc: &dplanemgrv1.DataPlane_Vpc{Cidr: "  "},
	}, nil).Once()

	action := &ApplyLoadBalanceAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbGroupRepo:   tidbGroupRepo,
		tidbClusterRepo: tidbClusterRepo,
		client:          client,
		dataplaneClient: client,
	}

	err = action.Invoke(actions.ActionArgs{
		workflow.UpdateLoadBalancerKey: actions.ActionArgString(workflow.UpdateLoadBalancerValue),
		workflow.SharedPoolIdKey:       actions.ActionArgString("shared-pool-1"),
		workflow.TiDBGroupIdKey:        actions.ActionArgString("tidb-group-1"),
		workflow.TiDBGroupKey:          actions.ActionArgString(string(tidbGroupStr)),
	})
	require.EqualError(t, err, "dataplane vpc cidr is empty")
}

func TestApplyLoadBalanceActionReturnsErrorWhenDataplaneCIDRInvalid(t *testing.T) {
	sharedPoolRepo := repomocks.NewSharedPoolRepo(t)
	tidbGroupRepo := repomocks.NewTiDBGroupRepo(t)
	tidbClusterRepo := repomocks.NewTiDBClusterRepo(t)
	client := dataplanemgrmocks.NewClient(t)

	sharedPool := &entity.SharedPool{
		SharedPool: apiv1.SharedPool{
			Id:            "shared-pool-1",
			DataplaneName: "dataplane-1",
			ComponentGroups: []*entity.Component{
				{
					Id:            "cluster-1",
					ComponentType: apiv1.ComponentType_TIPROXY,
				},
			},
		},
	}
	tidbGroup := &entity.TiDBGroup{
		TiDBGroup: apiv1.TiDBGroup{
			Id:            "tidb-group-1",
			ClusterId:     "cluster-1",
			SharedPoolId:  "shared-pool-1",
			DataplaneName: "dataplane-1",
			SqlEndpoint: &entity.SQLEndpoint{
				PublicLb: &entity.LoadBalancer{Enabled: false},
				PrivateLb: &entity.LoadBalancer{
					Enabled: true,
					Port:    4000,
				},
			},
		},
	}
	tidbCluster := &entity.TiDBCluster{
		Id:              "cluster-1",
		GatewayIdentity: "cluster-gateway",
	}

	tidbGroupStr, err := protojson.Marshal(tidbGroup)
	require.NoError(t, err)

	sharedPoolRepo.On("Load", mock.Anything, "shared-pool-1", false).Return(sharedPool, nil).Once()
	tidbGroupRepo.On("Load", mock.Anything, "tidb-group-1", false).Return(tidbGroup, nil).Once()
	tidbClusterRepo.On("Load", mock.Anything, "cluster-1", false).Return(tidbCluster, nil).Once()
	client.On("GetDataPlane", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.GetDataPlaneRequest) bool {
		return req.GetDataPlaneId() == "dataplane-1"
	})).Return(&dplanemgrv1.DataPlane{
		Vpc: &dplanemgrv1.DataPlane_Vpc{Cidr: "not-a-cidr"},
	}, nil).Once()

	action := &ApplyLoadBalanceAction{
		sharedPoolRepo:  sharedPoolRepo,
		tidbGroupRepo:   tidbGroupRepo,
		tidbClusterRepo: tidbClusterRepo,
		client:          client,
		dataplaneClient: client,
	}

	err = action.Invoke(actions.ActionArgs{
		workflow.UpdateLoadBalancerKey: actions.ActionArgString(workflow.UpdateLoadBalancerValue),
		workflow.SharedPoolIdKey:       actions.ActionArgString("shared-pool-1"),
		workflow.TiDBGroupIdKey:        actions.ActionArgString("tidb-group-1"),
		workflow.TiDBGroupKey:          actions.ActionArgString(string(tidbGroupStr)),
	})
	require.ErrorContains(t, err, "dataplane vpc cidr not-a-cidr is invalid")
}