go-migrate/internal/app/etl/table_analyzers/main_test.go

package table_analyzers

import (
	"context"
	"testing"

	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/config"
	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/etl"
	"git.ksdemosapps.com/kylesoda/go-migrate/internal/app/models"
)

type MockTableAnalyzer struct {
	minValue  int64
	maxValue  int64
	totalRows int64
	capturedRangeConstraint config.RangeConfig
}

func (m *MockTableAnalyzer) QueryColumnTypes(_ context.Context, _ config.TableInfo) ([]models.ColumnType, error) {
	return nil, nil
}

func (m *MockTableAnalyzer) EstimateTotalRows(_ context.Context, _ config.TableInfo) (int64, error) {
	return m.totalRows, nil
}

func (m *MockTableAnalyzer) QueryMaxMinFromColumn(_ context.Context, _ config.TableInfo, _ string) (etl.MaxMinColumnResult, error) {
	return etl.MaxMinColumnResult{Min: m.minValue, Max: m.maxValue}, nil
}

func (m *MockTableAnalyzer) CalculatePartitionRanges(_ context.Context, _ config.TableInfo, _ string, _ int64, rangeConstraint config.RangeConfig) ([]models.Partition, error) {
	m.capturedRangeConstraint = rangeConstraint
	return []models.Partition{}, nil
}

//go:fix inline
func ptr64(v int64) *int64 { return new(v) }

var testTableInfo = config.TableInfo{Schema: "dbo", Table: "test"}

func TestCalculatePartitionsEstimation_NoOverlap(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{minValue: 0, maxValue: 100}

	partitions, err := calculatePartitionsEstimation(ctx, mock, testTableInfo, "id", 4, config.RangeConfig{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(partitions) != 4 {
		t.Errorf("expected 4 partitions, got %d", len(partitions))
	}

	for i := 0; i < len(partitions)-1; i++ {
		current := partitions[i].Range
		next := partitions[i+1].Range
		if current.Max == next.Min && current.IsMaxInclusive && next.IsMinInclusive {
			t.Errorf("partition %d and %d overlap at value %d (both inclusive)", i, i+1, current.Max)
		}
	}
}

func TestCalculatePartitionsEstimation_CoverageComplete(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{minValue: 1000, maxValue: 2000}

	partitions, err := calculatePartitionsEstimation(ctx, mock, testTableInfo, "id", 5, config.RangeConfig{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if partitions[0].Range.Min != 1000 || !partitions[0].Range.IsMinInclusive {
		t.Errorf("first partition should start at 1000 (inclusive), got %d (inclusive=%v)",
			partitions[0].Range.Min, partitions[0].Range.IsMinInclusive)
	}

	if partitions[len(partitions)-1].Range.Max != 2000 {
		t.Errorf("last partition should end at 2000, got %d", partitions[len(partitions)-1].Range.Max)
	}
}

func TestCalculatePartitionsEstimation_FirstPartitionInclusive(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{minValue: 50, maxValue: 70}

	partitions, err := calculatePartitionsEstimation(ctx, mock, testTableInfo, "id", 3, config.RangeConfig{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if !partitions[0].Range.IsMinInclusive {
		t.Errorf("first partition should have IsMinInclusive=true")
	}

	if partitions[0].Range.Min != 50 {
		t.Errorf("first partition should start at 50, got %d", partitions[0].Range.Min)
	}

	for i := 1; i < len(partitions); i++ {
		if partitions[i].Range.IsMinInclusive {
			t.Errorf("partition %d should have IsMinInclusive=false to avoid overlap", i)
		}
	}
}

func TestPartitionRangeGenerator_Exact_NoRange_PassesEmptyConstraint(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 1000}

	_, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, config.RangeConfig{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if mock.capturedRangeConstraint.Min != nil || mock.capturedRangeConstraint.Max != nil {
		t.Errorf("expected empty range constraint, got min=%v max=%v",
			mock.capturedRangeConstraint.Min, mock.capturedRangeConstraint.Max)
	}
}

func TestPartitionRangeGenerator_Exact_BothBounds_PassesBothToAnalyzer(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 1000}
	jobRange := config.RangeConfig{Min: ptr64(200), Max: ptr64(800), IsMinInclusive: true, IsMaxInclusive: true}

	_, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rc := mock.capturedRangeConstraint
	if rc.Min == nil || *rc.Min != 200 {
		t.Errorf("expected Min=200, got %v", rc.Min)
	}
	if rc.Max == nil || *rc.Max != 800 {
		t.Errorf("expected Max=800, got %v", rc.Max)
	}
	if !rc.IsMinInclusive || !rc.IsMaxInclusive {
		t.Errorf("expected both bounds inclusive, got minInc=%v maxInc=%v", rc.IsMinInclusive, rc.IsMaxInclusive)
	}
}

func TestPartitionRangeGenerator_Exact_MinOnly_PassesMinNilMax(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 1000}
	jobRange := config.RangeConfig{Min: ptr64(500)}

	_, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rc := mock.capturedRangeConstraint
	if rc.Min == nil || *rc.Min != 500 {
		t.Errorf("expected Min=500, got %v", rc.Min)
	}
	if rc.Max != nil {
		t.Errorf("expected Max=nil (no upper bound), got %v", rc.Max)
	}
}

func TestPartitionRangeGenerator_Exact_MaxOnly_PassesMaxNilMin(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 1000}
	jobRange := config.RangeConfig{Max: ptr64(300)}

	_, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	rc := mock.capturedRangeConstraint
	if rc.Min != nil {
		t.Errorf("expected Min=nil (no lower bound), got %v", rc.Min)
	}
	if rc.Max == nil || *rc.Max != 300 {
		t.Errorf("expected Max=300, got %v", rc.Max)
	}
}

func TestPartitionRangeGenerator_Estimation_BothBounds_UsesUserRange(t *testing.T) {
	ctx := context.Background()
	// DB min/max differ intentionally — user bounds should take precedence.
	mock := &MockTableAnalyzer{totalRows: 1000, minValue: 0, maxValue: 999}
	jobRange := config.RangeConfig{Min: ptr64(200), Max: ptr64(700), IsMinInclusive: true}

	partitions, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "ESTIMATION", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(partitions) == 0 {
		t.Fatal("expected at least one partition")
	}

	if partitions[0].Range.Min != 200 {
		t.Errorf("first partition should start at user min=200, got %d", partitions[0].Range.Min)
	}
	if partitions[len(partitions)-1].Range.Max != 700 {
		t.Errorf("last partition should end at user max=700, got %d", partitions[len(partitions)-1].Range.Max)
	}
}

func TestPartitionRangeGenerator_Estimation_MinOnly_QueriesDBForMax(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 1000, minValue: 0, maxValue: 999}
	jobRange := config.RangeConfig{Min: ptr64(500), IsMinInclusive: true}

	partitions, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "ESTIMATION", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(partitions) == 0 {
		t.Fatal("expected at least one partition")
	}

	if partitions[0].Range.Min != 500 {
		t.Errorf("first partition should start at user min=500, got %d", partitions[0].Range.Min)
	}
	if partitions[len(partitions)-1].Range.Max != 999 {
		t.Errorf("last partition should end at DB max=999, got %d", partitions[len(partitions)-1].Range.Max)
	}
}

func TestPartitionRangeGenerator_Estimation_MaxOnly_QueriesDBForMin(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 1000, minValue: 100, maxValue: 999}
	jobRange := config.RangeConfig{Max: ptr64(600), IsMaxInclusive: true}

	partitions, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "ESTIMATION", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(partitions) == 0 {
		t.Fatal("expected at least one partition")
	}

	if partitions[0].Range.Min != 100 {
		t.Errorf("first partition should start at DB min=100, got %d", partitions[0].Range.Min)
	}
	if partitions[len(partitions)-1].Range.Max != 600 {
		t.Errorf("last partition should end at user max=600, got %d", partitions[len(partitions)-1].Range.Max)
	}
}

func TestPartitionRangeGenerator_SinglePartition_NoRange(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 50}

	partitions, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, config.RangeConfig{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(partitions) != 1 {
		t.Fatalf("expected 1 partition, got %d", len(partitions))
	}
	if partitions[0].HasRange {
		t.Error("single partition with no range should have HasRange=false")
	}
}

func TestPartitionRangeGenerator_SinglePartition_BothBounds(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 50}
	jobRange := config.RangeConfig{Min: ptr64(100), Max: ptr64(200), IsMinInclusive: true, IsMaxInclusive: true}

	partitions, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(partitions) != 1 {
		t.Fatalf("expected 1 partition, got %d", len(partitions))
	}
	p := partitions[0]
	if !p.HasRange {
		t.Error("expected HasRange=true")
	}
	if p.Range.Min != 100 || p.Range.Max != 200 {
		t.Errorf("expected [100, 200], got [%d, %d]", p.Range.Min, p.Range.Max)
	}
	if !p.Range.IsMinInclusive || !p.Range.IsMaxInclusive {
		t.Errorf("expected both inclusive, got minInc=%v maxInc=%v", p.Range.IsMinInclusive, p.Range.IsMaxInclusive)
	}
}

func TestPartitionRangeGenerator_SinglePartition_MinOnly(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 50}
	jobRange := config.RangeConfig{Min: ptr64(100), IsMinInclusive: true}

	partitions, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	p := partitions[0]
	if !p.HasRange {
		t.Error("expected HasRange=true")
	}
	if p.Range.Min != 100 {
		t.Errorf("expected Min=100, got %d", p.Range.Min)
	}
	if p.Range.Max != 0 {
		t.Errorf("expected Max=0 (no upper bound), got %d", p.Range.Max)
	}
}

func TestPartitionRangeGenerator_SinglePartition_MaxOnly(t *testing.T) {
	ctx := context.Background()
	mock := &MockTableAnalyzer{totalRows: 50}
	jobRange := config.RangeConfig{Max: ptr64(200), IsMaxInclusive: true}

	partitions, err := PartitionRangeGenerator(ctx, mock, testTableInfo, "id", "EXACT", 100, jobRange)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	p := partitions[0]
	if !p.HasRange {
		t.Error("expected HasRange=true")
	}
	if p.Range.Min != 0 {
		t.Errorf("expected Min=0 (no lower bound), got %d", p.Range.Min)
	}
	if p.Range.Max != 200 {
		t.Errorf("expected Max=200, got %d", p.Range.Max)
	}
}