package server

import (
	"context"
	"maps"
	"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"
	"github.com/tidbcloud/cluster-service-ng/domain/workflow/dag"
	dataplanemgrmocks "github.com/tidbcloud/cluster-service-ng/pkg/dataplanemgr/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/nodegroup"
	"github.com/tidbcloud/cluster-service-ng/pkg/runtime/actions"
	executormocks "github.com/tidbcloud/cluster-service-ng/pkg/runtime/executor/mocks"
	"github.com/tidbcloud/cluster-service-ng/pkg/storage/repo"
	repomnocks "github.com/tidbcloud/cluster-service-ng/pkg/storage/repo/mocks"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

func TestCreateNodeInstanceGroup(t *testing.T) {
	ctx := context.Background()

	tests := []struct {
		name        string
		group       *entity.NodeInstanceGroup
		setupMocks  func(*repomnocks.NodeInstanceGroupRepo)
		expectCode  codes.Code
		expectError bool
	}{
		{
			name:        "rejects nil group",
			group:       nil,
			setupMocks:  func(r *repomnocks.NodeInstanceGroupRepo) {},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name: "rejects missing primary instance type",
			group: &entity.NodeInstanceGroup{
				NodeSpec:            &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16},
				PrimaryInstanceType: "",
			},
			setupMocks:  func(r *repomnocks.NodeInstanceGroupRepo) {},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name: "rejects duplicate group",
			group: &entity.NodeInstanceGroup{
				NodeSpec:            &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16},
				PrimaryInstanceType: "m5.2xlarge",
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo) {
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(&entity.NodeInstanceGroup{}, nil).Once()
			},
			expectCode:  codes.AlreadyExists,
			expectError: true,
		},
		{
			name: "rejects fallback spec smaller than primary",
			group: &entity.NodeInstanceGroup{
				NodeSpec:            &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16},
				PrimaryInstanceType: "m5.2xlarge",
				FallbackInstances: []*apiv1.NodeInstanceGroup_FallbackInstance{
					{InstanceType: "m5.large", NodeSpec: &apiv1.NodeSpec{Cpu: 4, MemoryGi: 8}, Priority: 100},
				},
			},
			setupMocks:  func(r *repomnocks.NodeInstanceGroupRepo) {},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name: "creates group with PENDING status",
			group: &entity.NodeInstanceGroup{
				NodeSpec:            &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16},
				PrimaryInstanceType: "m5.2xlarge",
				FallbackInstances: []*apiv1.NodeInstanceGroup_FallbackInstance{
					{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
				},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo) {
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
				r.On("Save", mock.Anything, mock.MatchedBy(func(g *entity.NodeInstanceGroup) bool {
					return g.GetStatus() == apiv1.NodeInstanceGroup_PENDING &&
						g.GetPrimaryInstanceType() == "m5.2xlarge" &&
						g.GetName() == "8C16G"
				})).Return(nil).Once()
			},
			expectError: false,
		},
		{
			name: "rejects fallback missing priority",
			group: &entity.NodeInstanceGroup{
				NodeSpec:            &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16},
				PrimaryInstanceType: "m5.2xlarge",
				FallbackInstances: []*apiv1.NodeInstanceGroup_FallbackInstance{
					{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}},
				},
			},
			setupMocks:  func(r *repomnocks.NodeInstanceGroupRepo) {},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			nodeGroupRepo := repomnocks.NewNodeInstanceGroupRepo(t)
			tt.setupMocks(nodeGroupRepo)

			s := &serverImpl{
				nodeGroupRepo: nodeGroupRepo,
				transactioner: &mockTransactioner{},
			}
			result, err := s.CreateNodeInstanceGroup(ctx, tt.group)
			if tt.expectError {
				require.Error(t, err)
				require.Equal(t, tt.expectCode, status.Code(err))
				return
			}
			require.NoError(t, err)
			require.NotNil(t, result)
			require.Equal(t, apiv1.NodeInstanceGroup_PENDING, result.GetStatus())
		})
	}
}

func TestBuildNodeInstanceGroupDagArgs_PropagatesPlacementGroupEnabled(t *testing.T) {
	// The CSNG producer side computes placement_group_enabled from the
	// dataplane's node-zone count and threads it through the DAG args so the
	// downstream CreateNodePool action can include it in the gRPC request.
	// The DPM service still derives the same value from the VPC topology, so
	// the wire field stays a single source of truth.
	args, err := buildNodeInstanceGroupDagArgs(context.Background(), &entity.NodeInstanceGroup{
		NodeSpec:            &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16},
		PrimaryInstanceType: "m6i.2xlarge",
	}, "dp-1", true)
	require.NoError(t, err)

	enabled, err := args.GetBool(workflow.PlacementGroupEnabledKey)
	require.NoError(t, err)
	require.True(t, enabled)

	pbArgs, err := actions.BuildPBFromActionArgs(args)
	require.NoError(t, err)
	require.Equal(t, "true", pbArgs[workflow.PlacementGroupEnabledKey].GetActionArgString())
}

func TestPlacementGroupEnabledForDataplanesReadsSharedPoolGate(t *testing.T) {
	// The shared pool's Provider.PlacementGroupEnabled is the durable
	// per-pool gate, set at CreateSharedPool time from topology. Every
	// downstream consumer (NodePool infra, component templates, NIG
	// creation, shortage fallback) must read it instead of re-deriving
	// from the node-zone list, otherwise a legacy single-AZ/PG-disabled
	// pool ends up with PG-disabled infra but PG-enabled scheduling.
	ctx := context.Background()
	dataplanes := []string{"dp-1", "dp-legacy", "dp-mz"}
	tests := []struct {
		name     string
		pools    []*entity.SharedPool
		expected map[string]bool
	}{
		{
			name: "shared pool with gate=true is eligible",
			pools: []*entity.SharedPool{
				{SharedPool: apiv1.SharedPool{DataplaneName: "dp-1", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a", "us-west-2b", "us-west-2c"}, PlacementGroupEnabled: true}}},
			},
			expected: map[string]bool{"dp-1": true, "dp-legacy": false, "dp-mz": false},
		},
		{
			name: "shared pool with gate=false is not eligible even on single-AZ",
			pools: []*entity.SharedPool{
				{SharedPool: apiv1.SharedPool{DataplaneName: "dp-legacy", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a"}, PlacementGroupEnabled: false}}},
			},
			expected: map[string]bool{"dp-1": false, "dp-legacy": false, "dp-mz": false},
		},
		{
			name:     "missing shared pool defaults to false",
			pools:    nil,
			expected: map[string]bool{"dp-1": false, "dp-legacy": false, "dp-mz": false},
		},
		{
			// PG eligibility follows the per-pool gate, not the live zone
			// count. A legacy single-AZ shared pool with gate=false stays
			// PG-disabled, and a future multi-AZ shared pool with gate=true
			// stays PG-enabled.
			name: "multi-AZ shared pool with gate=true is eligible",
			pools: []*entity.SharedPool{
				{SharedPool: apiv1.SharedPool{DataplaneName: "dp-mz", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a", "us-west-2b"}, PlacementGroupEnabled: true}}},
			},
			expected: map[string]bool{"dp-1": false, "dp-legacy": false, "dp-mz": true},
		},
		{
			// Repo ignores the DataplaneNames filter and returns a pool
			// for an unrelated dataplane. The result map must not be
			// widened: only requested dataplanes appear, and the
			// unrelated entry is silently dropped.
			name: "repo returning an unrelated pool does not widen the result map",
			pools: []*entity.SharedPool{
				{SharedPool: apiv1.SharedPool{DataplaneName: "dp-1", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a"}, PlacementGroupEnabled: true}}},
				{SharedPool: apiv1.SharedPool{DataplaneName: "dp-unexpected", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a"}, PlacementGroupEnabled: true}}},
			},
			expected: map[string]bool{"dp-1": true, "dp-legacy": false, "dp-mz": false},
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			sharedPoolRepo.On("List", mock.Anything, &repo.SharedPoolFilter{DataplaneNames: dataplanes}).Return(tt.pools, nil).Once()

			s := &serverImpl{sharedPoolRepo: sharedPoolRepo}
			got, err := s.placementGroupEnabledForDataplanes(ctx, dataplanes)
			require.NoError(t, err)
			require.Equal(t, tt.expected, got)
		})
	}
}

func TestPlacementGroupEnabledForDataplanesPropagatesRepoErrors(t *testing.T) {
	// SharedPool lookup errors must surface so the DAG retries rather than
	// silently submitting a wrong wire value. Old DPM builds still pin the
	// request value as the durable spec, so a swallowed error here would
	// pin false on the new pool and break producer/consumer rollout.
	ctx := context.Background()
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	sharedPoolRepo.On("List", mock.Anything, mock.Anything).Return(nil, assertAnError())

	s := &serverImpl{sharedPoolRepo: sharedPoolRepo}
	got, err := s.placementGroupEnabledForDataplanes(ctx, []string{"dp-1"})
	require.Error(t, err)
	require.Nil(t, got)
}

func TestPlacementGroupEnabledForDataplaneRequiresRepo(t *testing.T) {
	// No sharedPoolRepo means the call cannot honour the durable migration
	// contract; fail closed with an error instead of pinning false.
	got, err := (&serverImpl{}).placementGroupEnabledForDataplane(context.Background(), "dp-1")
	require.Error(t, err)
	require.False(t, got)
}

type errSentinel struct{}

func (errSentinel) Error() string { return "boom" }

func assertAnError() error { return errSentinel{} }

func TestSetNodeInstanceGroupInstanceTypeStatusDoesNotResolvePlacementGroupEnabled(t *testing.T) {
	// SetNodeInstanceGroupInstanceTypeStatus only taints the affected
	// instance types; it never creates NodePools and never reads
	// placement_group_enabled. The DAG args carry the literal false and
	// the SharedPool repo must not be queried — otherwise the call gains
	// a spurious failure mode for a value the action ignores.
	ctx := context.Background()
	nodeGroupRepo := repomnocks.NewNodeInstanceGroupRepo(t)
	dataplaneNodeGroupRepo := repomnocks.NewDataplaneNodeInstanceGroupRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	dpmClient := dataplanemgrmocks.NewClient(t)

	group := buildNodeInstanceGroup(nil)
	group.Status = apiv1.NodeInstanceGroup_ACTIVE
	group.FallbackInstances = []*apiv1.NodeInstanceGroup_FallbackInstance{
		{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
	}
	nodeGroupRepo.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(group, nil).Once()
	nodeGroupRepo.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
	dpmClient.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
		DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}},
	}, nil).Once()
	dataplaneNodeGroupRepo.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").Return(nil, nil).Once()
	// sharedPoolRepo.List is intentionally NOT mocked — calling it would
	// fail the test (the testify mock framework reports unstubbed calls).
	dagExecutor.On("Submit", mock.Anything, dag.PlanSetNodeInstanceGroupInstanceTypeStatus, mock.MatchedBy(func(args actions.ActionArgs) bool {
		value, err := args.GetBool(workflow.PlacementGroupEnabledKey)
		return err == nil && !value
	}), mock.Anything).Return(nil).Once()

	s := &serverImpl{
		nodeGroupRepo:          nodeGroupRepo,
		dataplaneNodeGroupRepo: dataplaneNodeGroupRepo,
		sharedPoolRepo:         sharedPoolRepo,
		dagExecutor:            dagExecutor,
		dataplaneClient:        dpmClient,
		transactioner:          &mockTransactioner{},
	}

	_, err := s.SetNodeInstanceGroupInstanceTypeStatus(ctx, "8C16G", "m5.2xlarge", "m5.4xlarge", apiv1.InstanceTypeStatus_DISABLED)
	require.NoError(t, err)
}

func TestSetNodeInstanceGroupStatusUsesPerDataplanePlacementGroupEnabled(t *testing.T) {
	// Each dataplane carries its own per-pool migration gate on
	// sharedPool.Provider.PlacementGroupEnabled. The DAG submit must carry
	// the right per-dataplane value so the downstream gRPC CreateNodePool
	// request can be routed to the correct DPM branch.
	ctx := context.Background()
	nodeGroupRepo := repomnocks.NewNodeInstanceGroupRepo(t)
	dataplaneNodeGroupRepo := repomnocks.NewDataplaneNodeInstanceGroupRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	dpmClient := dataplanemgrmocks.NewClient(t)

	group := buildNodeInstanceGroup(nil)
	group.Status = apiv1.NodeInstanceGroup_PENDING
	nodeGroupRepo.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(group, nil).Once()
	nodeGroupRepo.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
	dpmClient.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
		DataPlanes: []*dplanemgrv1.DataPlane{
			{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active},
			{Id: "dp-2", Status: dplanemgrv1.DataPlane_Active},
		},
	}, nil).Once()
	sharedPoolRepo.On("List", mock.Anything, mock.MatchedBy(func(f *repo.SharedPoolFilter) bool {
		return f != nil && len(f.DataplaneNames) == 1 && f.DataplaneNames[0] == "dp-1"
	})).Return([]*entity.SharedPool{
		// dp-1 is a new ZONAL pool: gate=true.
		{SharedPool: apiv1.SharedPool{DataplaneName: "dp-1", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a"}, PlacementGroupEnabled: true}}},
	}, nil).Once()
	sharedPoolRepo.On("List", mock.Anything, mock.MatchedBy(func(f *repo.SharedPoolFilter) bool {
		return f != nil && len(f.DataplaneNames) == 1 && f.DataplaneNames[0] == "dp-2"
	})).Return([]*entity.SharedPool{
		// dp-2 is a legacy REGIONAL pool opted out of PG: gate=false even
		// though its node-zone list happens to also be single-AZ.
		{SharedPool: apiv1.SharedPool{DataplaneName: "dp-2", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a", "us-west-2b", "us-west-2c"}, PlacementGroupEnabled: false}}},
	}, nil).Once()
	dataplaneNodeGroupRepo.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").Return(nil, nil).Once()
	dataplaneNodeGroupRepo.On("LoadByDataplaneAndSpec", mock.Anything, "dp-2", "8C16G", "m5.2xlarge").Return(nil, nil).Once()
	dagExecutor.On("Submit", mock.Anything, dag.PlanCreateNodeInstanceGroup, mock.MatchedBy(func(args actions.ActionArgs) bool {
		dataplaneName, err := args.GetString(workflow.DataplaneIdKey)
		if err != nil || dataplaneName != "dp-1" {
			return false
		}
		placementGroupEnabled, err := args.GetBool(workflow.PlacementGroupEnabledKey)
		return err == nil && placementGroupEnabled
	}), mock.Anything).Return(nil).Once()
	dagExecutor.On("Submit", mock.Anything, dag.PlanCreateNodeInstanceGroup, mock.MatchedBy(func(args actions.ActionArgs) bool {
		dataplaneName, err := args.GetString(workflow.DataplaneIdKey)
		if err != nil || dataplaneName != "dp-2" {
			return false
		}
		placementGroupEnabled, err := args.GetBool(workflow.PlacementGroupEnabledKey)
		return err == nil && !placementGroupEnabled
	}), mock.Anything).Return(nil).Once()

	s := &serverImpl{
		nodeGroupRepo:          nodeGroupRepo,
		dataplaneNodeGroupRepo: dataplaneNodeGroupRepo,
		sharedPoolRepo:         sharedPoolRepo,
		dagExecutor:            dagExecutor,
		dataplaneClient:        dpmClient,
		transactioner:          &mockTransactioner{},
	}

	_, err := s.SetNodeInstanceGroupStatus(ctx, "8C16G", "m5.2xlarge", apiv1.NodeInstanceGroup_ACTIVE)
	require.NoError(t, err)
}

func TestSetNodeInstanceGroupStatus(t *testing.T) {
	ctx := context.Background()

	tests := []struct {
		name            string
		nodeSpec        string
		primaryInstance string
		targetStatus    apiv1.NodeInstanceGroup_Status
		setupMocks      func(*repomnocks.NodeInstanceGroupRepo, *repomnocks.DataplaneNodeInstanceGroupRepo, *repomnocks.DataplaneNodeInstanceGroupStateRepo, *repomnocks.SharedPoolRepo, *executormocks.RemoteExecutor, *dataplanemgrmocks.Client)
		expectCode      codes.Code
		expectError     bool
	}{
		{
			name:            "rejects PENDING status",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_PENDING,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:            "allows DEPRECATED to ACTIVE transition",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_ACTIVE,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_DEPRECATED
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}},
				}, nil).Once()
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
				e.On("Submit", mock.Anything, dag.PlanReactivateNodeInstanceGroup, mock.Anything, mock.Anything).
					Return(nil).Once()
			},
			expectError: false,
		},
		{
			name:            "rejects unspecified status",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_STATUS_UNSPECIFIED,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:            "allows ACTIVE when another primary is already ACTIVE",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_ACTIVE,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_PENDING
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}},
				}, nil).Once()
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
				e.On("Submit", mock.Anything, dag.PlanCreateNodeInstanceGroup, mock.Anything, mock.Anything).
					Return(nil).Once()
			},
			expectError: false,
		},
		{
			name:            "sets ACTIVE and submits create DAGs",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_ACTIVE,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_PENDING
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}},
				}, nil).Once()
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
				e.On("Submit", mock.Anything, dag.PlanCreateNodeInstanceGroup, mock.Anything, mock.Anything).
					Return(nil).Once()
			},
			expectError: false,
		},
		{
			name:            "skips deleting dataplanes",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_ACTIVE,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_PENDING
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				// dp-2 is Deleting and should be filtered out by listActiveDataplaneNames.
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{
						{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active},
						{Id: "dp-2", Status: dplanemgrv1.DataPlane_Deleting},
					},
				}, nil).Once()
				// Only dp-1 should be processed.
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
				e.On("Submit", mock.Anything, dag.PlanCreateNodeInstanceGroup, mock.Anything, mock.Anything).
					Return(nil).Once()
			},
			expectError: false,
		},
		{
			name:            "skips dataplane with deprecated override",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_ACTIVE,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_PENDING
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}},
				}, nil).Once()
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(&entity.DataplaneNodeInstanceGroup{
						Status: apiv1.DataplaneNodeInstanceGroup_DEPRECATED,
					}, nil).Once()
			},
			expectError: false,
		},
		{
			name:            "deprecate skips dataplanes where group was never active",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_DEPRECATED,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				// Group is PENDING, so it was never active on dataplanes without an ACTIVE overlay.
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_PENDING
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}, {Id: "dp-2", Status: dplanemgrv1.DataPlane_Active}},
				}, nil).Once()
				// dp-1 has ACTIVE override → was active → but ACTIVE override means skip (override takes precedence)
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(&entity.DataplaneNodeInstanceGroup{
						Status: apiv1.DataplaneNodeInstanceGroup_ACTIVE,
					}, nil).Once()
				// dp-2 has no override → group was PENDING, never active → skip, no DAG submitted
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-2", "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
				// No DAG should be submitted for either dataplane.
			},
			expectError: false,
		},
		{
			name:            "deprecates active group and cleans up shortage detector states",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_DEPRECATED,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_ACTIVE
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}},
				}, nil).Once()
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
				e.On("Submit", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
					Return(nil).Once()
				// State cleanup should be called for the deprecated dataplane.
				st.On("DeleteByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(nil).Once()
			},
			expectError: false,
		},
		{
			name:            "rejects RETIRING from ACTIVE (must be DEPRECATED first)",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_RETIRING,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_ACTIVE
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{}, nil).Once()
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
			},
			expectCode:  codes.FailedPrecondition,
			expectError: true,
		},
		{
			name:            "rejects RETIRING when dataplane override is ACTIVE",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_RETIRING,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_DEPRECATED
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{}, nil).Once()
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				// One override is still ACTIVE, blocking the transition.
				d.On("ListBySpec", mock.Anything, "8C16G").Return([]*entity.DataplaneNodeInstanceGroup{
					{
						DataplaneName:       "dp-1",
						PrimaryInstanceType: "m5.2xlarge",
						Status:              apiv1.DataplaneNodeInstanceGroup_ACTIVE,
					},
				}, nil).Once()
			},
			expectCode:  codes.FailedPrecondition,
			expectError: true,
		},
		{
			name:            "sets RETIRING with bulk state cleanup",
			nodeSpec:        "8C16G",
			primaryInstance: "m5.2xlarge",
			targetStatus:    apiv1.NodeInstanceGroup_RETIRING,
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_DEPRECATED
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				// All overrides must be DEPRECATED for RETIRING precondition.
				d.On("ListBySpec", mock.Anything, "8C16G").Return([]*entity.DataplaneNodeInstanceGroup{
					{
						DataplaneName:       "dp-1",
						PrimaryInstanceType: "m5.2xlarge",
						Status:              apiv1.DataplaneNodeInstanceGroup_DEPRECATED,
					},
				}, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				c.On("ListDataPlanes", mock.Anything, mock.Anything).Return(&dplanemgrv1.ListDataPlanesResponse{
					DataPlanes: []*dplanemgrv1.DataPlane{{Id: "dp-1", Status: dplanemgrv1.DataPlane_Active}},
				}, nil).Once()
				d.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
					Return(&entity.DataplaneNodeInstanceGroup{
						DataplaneName:       "dp-1",
						PrimaryInstanceType: "m5.2xlarge",
						Status:              apiv1.DataplaneNodeInstanceGroup_DEPRECATED,
					}, nil).Once()
				d.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
				// Bulk state cleanup should be called after RETIRING transition.
				st.On("DeleteBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(nil).Once()
			},
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			nodeGroupRepo := repomnocks.NewNodeInstanceGroupRepo(t)
			dpNodeGroupRepo := repomnocks.NewDataplaneNodeInstanceGroupRepo(t)
			stateRepo := repomnocks.NewDataplaneNodeInstanceGroupStateRepo(t)
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			// Default: no shared pool, so every dataplane gets gate=false.
			// Subtests that exercise the gate override this in setupMocks.
			sharedPoolRepo.On("List", mock.Anything, mock.Anything).Return([]*entity.SharedPool{}, nil).Maybe()
			dagExecutor := executormocks.NewRemoteExecutor(t)
			dpmClient := dataplanemgrmocks.NewClient(t)
			tt.setupMocks(nodeGroupRepo, dpNodeGroupRepo, stateRepo, sharedPoolRepo, dagExecutor, dpmClient)

			s := &serverImpl{
				nodeGroupRepo:          nodeGroupRepo,
				dataplaneNodeGroupRepo: dpNodeGroupRepo,
				nodeGroupStateRepo:     stateRepo,
				sharedPoolRepo:         sharedPoolRepo,
				dagExecutor:            dagExecutor,
				dataplaneClient:        dpmClient,
				transactioner:          &mockTransactioner{},
			}
			result, err := s.SetNodeInstanceGroupStatus(ctx, tt.nodeSpec, tt.primaryInstance, tt.targetStatus)
			if tt.expectError {
				require.Error(t, err)
				require.Equal(t, tt.expectCode, status.Code(err))
				return
			}
			require.NoError(t, err)
			require.NotNil(t, result)
		})
	}
}

func TestDeleteNodeInstanceGroup(t *testing.T) {
	ctx := context.Background()

	tests := []struct {
		name        string
		nodeSpec    string
		primary     string
		setupMocks  func(*repomnocks.NodeInstanceGroupRepo, *repomnocks.DataplaneNodeInstanceGroupRepo, *repomnocks.DataplaneNodeInstanceGroupStateRepo, *repomnocks.SharedPoolRepo, *executormocks.RemoteExecutor, *dataplanemgrmocks.Client)
		expectCode  codes.Code
		expectError bool
	}{
		{
			name:     "rejects when target group is not retired (ACTIVE)",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				// Target group is still ACTIVE → should reject
				target := buildNodeInstanceGroup(nil)
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(target, nil).Once()
			},
			expectCode:  codes.FailedPrecondition,
			expectError: true,
		},
		{
			name:     "rejects when target group is deprecated but not retired",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				target := buildNodeInstanceGroup(nil)
				target.Status = apiv1.NodeInstanceGroup_DEPRECATED
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(target, nil).Once()
			},
			expectCode:  codes.FailedPrecondition,
			expectError: true,
		},
		{
			name:     "rejects when target group is retiring but not retired",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				target := buildNodeInstanceGroup(nil)
				target.Status = apiv1.NodeInstanceGroup_RETIRING
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(target, nil).Once()
			},
			expectCode:  codes.FailedPrecondition,
			expectError: true,
		},
		{
			name:     "deletes retired group successfully",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, st *repomnocks.DataplaneNodeInstanceGroupStateRepo, sp *repomnocks.SharedPoolRepo, e *executormocks.RemoteExecutor, c *dataplanemgrmocks.Client) {
				target := buildNodeInstanceGroup(nil)
				target.Status = apiv1.NodeInstanceGroup_RETIRED
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(target, nil).Once()
				d.On("ListBySpec", mock.Anything, "8C16G").Return([]*entity.DataplaneNodeInstanceGroup{}, nil).Once()
				r.On("Delete", mock.Anything, "8C16G", "m5.2xlarge").Return(nil).Once()
				st.On("DeleteBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(nil).Once()
			},
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			nodeGroupRepo := repomnocks.NewNodeInstanceGroupRepo(t)
			dpNodeGroupRepo := repomnocks.NewDataplaneNodeInstanceGroupRepo(t)
			stateRepo := repomnocks.NewDataplaneNodeInstanceGroupStateRepo(t)
			sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
			sharedPoolRepo.On("List", mock.Anything, mock.Anything).Return([]*entity.SharedPool{}, nil).Maybe()
			dagExecutor := executormocks.NewRemoteExecutor(t)
			dpmClient := dataplanemgrmocks.NewClient(t)
			tt.setupMocks(nodeGroupRepo, dpNodeGroupRepo, stateRepo, sharedPoolRepo, dagExecutor, dpmClient)

			s := &serverImpl{
				nodeGroupRepo:          nodeGroupRepo,
				dataplaneNodeGroupRepo: dpNodeGroupRepo,
				nodeGroupStateRepo:     stateRepo,
				dagExecutor:            dagExecutor,
				dataplaneClient:        dpmClient,
				transactioner:          &mockTransactioner{},
			}
			_, err := s.DeleteNodeInstanceGroup(ctx, tt.nodeSpec, tt.primary)
			if tt.expectError {
				require.Error(t, err)
				require.Equal(t, tt.expectCode, status.Code(err))
				return
			}
			require.NoError(t, err)
		})
	}
}

func TestUpdateNodeInstanceGroup(t *testing.T) {
	ctx := context.Background()

	tests := []struct {
		name        string
		nodeSpec    string
		primary     string
		fallbacks   []*apiv1.NodeInstanceGroup_FallbackInstance
		setupMocks  func(*repomnocks.NodeInstanceGroupRepo, *repomnocks.DataplaneNodeInstanceGroupRepo, *dataplanemgrmocks.Client, *executormocks.RemoteExecutor)
		expectCode  codes.Code
		expectError bool
	}{
		{
			name:     "rejects empty node_spec",
			nodeSpec: "",
			primary:  "m5.2xlarge",
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "rejects not found group",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(nil, nil).Once()
			},
			expectCode:  codes.NotFound,
			expectError: true,
		},
		{
			name:     "rejects fallback same as primary",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.2xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16}, Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "rejects fallback missing priority",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "rejects fallback with nil node_spec",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.4xlarge", Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "rejects duplicate fallback instance",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
				{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 200},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "rejects fallback spec smaller than primary",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.large", NodeSpec: &apiv1.NodeSpec{Cpu: 4, MemoryGi: 8}, Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(buildNodeInstanceGroup(nil), nil).Once()
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "rejects fallback with local disk group",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
				group := buildNodeInstanceGroup(nil)
				group.NodeSpec.WithLocalDisk = true
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "rejects removal of existing fallback",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5a.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
				existing := buildNodeInstanceGroup([]*apiv1.NodeInstanceGroup_FallbackInstance{
					{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
				})
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(existing, nil).Once()
			},
			expectCode:  codes.InvalidArgument,
			expectError: true,
		},
		{
			name:     "updates fallback instances successfully",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(buildNodeInstanceGroup(nil), nil).Once()
				r.On("Save", mock.Anything, mock.MatchedBy(func(g *entity.NodeInstanceGroup) bool {
					return len(g.GetFallbackInstances()) == 1 &&
						g.GetFallbackInstances()[0].GetInstanceType() == "m5.4xlarge"
				})).Return(nil).Once()
			},
			expectError: false,
		},
		{
			name:     "updates fallback for pending group",
			nodeSpec: "8C16G",
			primary:  "m5.2xlarge",
			fallbacks: []*apiv1.NodeInstanceGroup_FallbackInstance{
				{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
			},
			setupMocks: func(r *repomnocks.NodeInstanceGroupRepo, d *repomnocks.DataplaneNodeInstanceGroupRepo, c *dataplanemgrmocks.Client, e *executormocks.RemoteExecutor) {
				group := buildNodeInstanceGroup(nil)
				group.Status = apiv1.NodeInstanceGroup_PENDING
				r.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").
					Return(group, nil).Once()
				r.On("Save", mock.Anything, mock.Anything).Return(nil).Once()
			},
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			nodeGroupRepo := repomnocks.NewNodeInstanceGroupRepo(t)
			dpNodeGroupRepo := repomnocks.NewDataplaneNodeInstanceGroupRepo(t)
			dpmClient := dataplanemgrmocks.NewClient(t)
			dagExecutor := executormocks.NewRemoteExecutor(t)
			tt.setupMocks(nodeGroupRepo, dpNodeGroupRepo, dpmClient, dagExecutor)

			s := &serverImpl{
				nodeGroupRepo:          nodeGroupRepo,
				dataplaneNodeGroupRepo: dpNodeGroupRepo,
				dataplaneClient:        dpmClient,
				dagExecutor:            dagExecutor,
				transactioner:          &mockTransactioner{},
			}
			result, err := s.UpdateNodeInstanceGroup(ctx, tt.nodeSpec, tt.primary, tt.fallbacks)
			if tt.expectError {
				require.Error(t, err)
				require.Equal(t, tt.expectCode, status.Code(err))
				return
			}
			require.NoError(t, err)
			require.NotNil(t, result)
		})
	}
}

func TestSetDataplaneNodeInstanceGroupStatusActiveUsesReactivateForDeprecated(t *testing.T) {
	ctx := context.Background()
	nodeGroupRepo := repomnocks.NewNodeInstanceGroupRepo(t)
	dataplaneNodeGroupRepo := repomnocks.NewDataplaneNodeInstanceGroupRepo(t)
	sharedPoolRepo := repomnocks.NewSharedPoolRepo(t)
	dagExecutor := executormocks.NewRemoteExecutor(t)
	dpmClient := dataplanemgrmocks.NewClient(t)

	group := buildNodeInstanceGroup(nil)
	group.Status = apiv1.NodeInstanceGroup_DEPRECATED
	nodeGroupRepo.On("LoadBySpecAndPrimary", mock.Anything, "8C16G", "m5.2xlarge").Return(group, nil).Once()
	dpmClient.On("GetDataPlane", mock.Anything, &dplanemgrv1.GetDataPlaneRequest{DataPlaneId: "dp-1"}).
		Return(&dplanemgrv1.DataPlane{Status: dplanemgrv1.DataPlane_Active}, nil).Once()
	dataplaneNodeGroupRepo.On("LoadByDataplaneAndSpec", mock.Anything, "dp-1", "8C16G", "m5.2xlarge").
		Return(&entity.DataplaneNodeInstanceGroup{
			DataplaneName:       "dp-1",
			NodeSpec:            group.GetNodeSpec(),
			PrimaryInstanceType: "m5.2xlarge",
			Status:              apiv1.DataplaneNodeInstanceGroup_DEPRECATED,
		}, nil).Once()
	dataplaneNodeGroupRepo.On("Save", mock.Anything, mock.MatchedBy(func(group *entity.DataplaneNodeInstanceGroup) bool {
		return group.GetStatus() == apiv1.DataplaneNodeInstanceGroup_ACTIVE
	})).Return(nil).Once()
	sharedPoolRepo.On("List", mock.Anything, &repo.SharedPoolFilter{DataplaneNames: []string{"dp-1"}}).Return([]*entity.SharedPool{
		{SharedPool: apiv1.SharedPool{DataplaneName: "dp-1", Provider: &apiv1.Provider{NodeZones: []string{"us-west-2a"}, PlacementGroupEnabled: true}}},
	}, nil).Once()
	dagExecutor.On("Submit", mock.Anything, dag.PlanReactivateNodeInstanceGroup, mock.Anything, mock.Anything).Return(nil).Once()

	s := &serverImpl{
		nodeGroupRepo:          nodeGroupRepo,
		dataplaneNodeGroupRepo: dataplaneNodeGroupRepo,
		sharedPoolRepo:         sharedPoolRepo,
		dagExecutor:            dagExecutor,
		dataplaneClient:        dpmClient,
		transactioner:          &mockTransactioner{},
	}

	result, err := s.SetDataplaneNodeInstanceGroupStatus(ctx, "dp-1", "8C16G", "m5.2xlarge", apiv1.DataplaneNodeInstanceGroup_ACTIVE)
	require.NoError(t, err)
	require.Equal(t, apiv1.DataplaneNodeInstanceGroup_ACTIVE, result.GetStatus())
}

func TestEnsureNoActiveNodePoolsForGroup(t *testing.T) {
	ctx := context.Background()

	tests := []struct {
		name        string
		group       *entity.NodeInstanceGroup
		nodePools   []*dplanemgrv1.NodePool
		expectCode  codes.Code
		expectError bool
	}{
		{
			name: "rejects fallback with replicas",
			group: buildNodeInstanceGroup([]*entity.NodeInstanceGroupFallbackInstance{
				{
					InstanceType: "m5.4xlarge",
					NodeSpec:     &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32},
					Priority:     100,
				},
			}),
			nodePools: []*dplanemgrv1.NodePool{
				makeNodePool("m5.2xlarge", "8C16G", 0, true, nil),
				makeNodePool("m5.4xlarge", "8C16G", 1, true, nil),
			},
			expectCode:  codes.FailedPrecondition,
			expectError: true,
		},
		{
			name: "ignores nodepool outside current fallback list",
			group: buildNodeInstanceGroup([]*entity.NodeInstanceGroupFallbackInstance{
				{
					InstanceType: "m5.4xlarge",
					NodeSpec:     &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32},
					Priority:     100,
				},
			}),
			nodePools: []*dplanemgrv1.NodePool{
				makeNodePool("m5.2xlarge", "8C16G", 0, true, nil),
				makeNodePool("m5.24xlarge", "8C16G", 1, false, map[string]string{
					nodegroup.LabelRealNodeSpec: "64c128g",
				}),
				makeNodePool("m5.8xlarge", "4C8G", 2, false, nil),
			},
			expectCode:  codes.OK,
			expectError: false,
		},
		{
			name:  "ignores fallback nodepool with same instance type as primary group",
			group: buildNodeInstanceGroup(nil),
			nodePools: []*dplanemgrv1.NodePool{
				makeNodePool("m5.2xlarge", "8C16G", 0, true, map[string]string{
					nodegroup.LabelPriority: "1000",
				}),
				makeNodePool("m5.2xlarge", "8C16G", 1, false, map[string]string{
					nodegroup.LabelPriority: "100",
				}),
			},
			expectCode:  codes.OK,
			expectError: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			client := dataplanemgrmocks.NewClient(t)
			client.On("ListNodePools", mock.Anything, mock.MatchedBy(func(req *dplanemgrv1.ListNodePoolsRequest) bool {
				return req.GetDataPlaneId() == "dp-1"
			})).Return(&dplanemgrv1.ListNodePoolsResponse{
				NodePools: tt.nodePools,
			}, nil).Once()

			s := &serverImpl{
				dataplaneClient: client,
			}

			err := s.ensureNoActiveNodePoolsForGroup(ctx, tt.group, []string{"dp-1"})
			if tt.expectError {
				require.Error(t, err)
				require.Equal(t, tt.expectCode, status.Code(err))
				return
			}
			require.NoError(t, err)
		})
	}
}

func buildNodeInstanceGroup(fallbacks []*entity.NodeInstanceGroupFallbackInstance) *entity.NodeInstanceGroup {
	nodeSpec := "8C16G"
	cpu, mem, _ := nodegroup.ParseSpec(nodeSpec)
	group := &entity.NodeInstanceGroup{
		Name:                nodeSpec,
		NodeSpec:            &apiv1.NodeSpec{Cpu: cpu, MemoryGi: mem},
		PrimaryInstanceType: "m5.2xlarge",
		Status:              apiv1.NodeInstanceGroup_ACTIVE,
	}
	group.FallbackInstances = fallbacks
	return group
}

func TestNodeInstanceGroupInstanceTypesKeepsPrimaryRole(t *testing.T) {
	group := buildNodeInstanceGroup([]*entity.NodeInstanceGroupFallbackInstance{
		{InstanceType: "m5.2xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 8, MemoryGi: 16}, Priority: 100},
		{InstanceType: "m5.4xlarge", NodeSpec: &apiv1.NodeSpec{Cpu: 16, MemoryGi: 32}, Priority: 100},
	})

	types := nodeInstanceGroupInstanceTypes(group)

	require.Equal(t, nodegroup.NodePoolRolePrimary, types["m5.2xlarge"])
	require.Equal(t, nodegroup.NodePoolRoleFallback, types["m5.4xlarge"])
}

func makeNodePool(instanceType, nodeSpec string, replicas uint32, deprecated bool, extraLabels map[string]string) *dplanemgrv1.NodePool {
	labels := map[string]string{
		nodegroup.LabelNodeSpec:  nodeSpec,
		nodegroup.LabelManagedBy: nodegroup.ManagedByInstanceTypeManager,
	}
	maps.Copy(labels, extraLabels)
	var taints []*dplanemgrv1.NodePoolTaint
	if deprecated {
		taint := nodegroup.DeprecatedTaint()
		taints = []*dplanemgrv1.NodePoolTaint{
			{
				Key:    taint.Key,
				Value:  taint.Value,
				Effect: taint.Effect,
			},
		}
	}
	return &dplanemgrv1.NodePool{
		InstanceType:    instanceType,
		Labels:          labels,
		Taints:          taints,
		CurrentReplicas: replicas,
	}
}