stillbox/pkg/database/partman/partman_test.go

package partman_test

import (
	"context"
	"testing"
	"time"

	"dynatron.me/x/stillbox/internal/common"
	"dynatron.me/x/stillbox/pkg/config"
	"dynatron.me/x/stillbox/pkg/database"
	"dynatron.me/x/stillbox/pkg/database/mocks"
	"dynatron.me/x/stillbox/pkg/database/partman"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
	"github.com/stretchr/testify/require"
)

var mctx = mock.Anything

func inTx(s *mocks.Store, dbtx *mocks.DBTX) {
	s.EXPECT().InTx(mctx, mock.AnythingOfType("func(database.Store) error"), mock.AnythingOfType("pgx.TxOptions")).RunAndReturn(func(ctx context.Context, f func(db database.Store) error, po pgx.TxOptions) error {
		return f(s)
	})

}

type timeRange struct {
	start time.Time
	end   time.Time
}

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

	timeInUTC := func(s string) time.Time {
		t, err := time.ParseInLocation("2006-01-02 15:04:05", s, time.UTC)
		if err != nil {
			panic(err)
		}

		return t
	}

	dateInUTC := func(s string) time.Time {
		t, err := time.ParseInLocation("2006-01-02", s, time.UTC)
		if err != nil {
			panic(err)
		}

		return t
	}

	tests := []struct {
		name          string
		now           time.Time
		cfg           config.Partition
		extant        []string
		expectCreate  []string
		expectDrop    []string
		expectDetach  []string
		expectSweep   []timeRange
		expectCleanup []timeRange
	}{
		{
			name: "monthly base",
			now:  timeInUTC("2024-11-28 11:37:04"),
			cfg: config.Partition{
				Enabled:      true,
				Schema:       "public",
				Interval:     "monthly",
				Retain:       2,
				Drop:         true,
				PreProvision: common.PtrTo(2),
			},
			extant: []string{
				"calls_p_2024_10",
				"calls_p_2024_09",
				"calls_p_2024_08",
				"calls_p_2024_07",
			},
			expectCreate: []string{
				"calls_p_2024_11",
				"calls_p_2024_12",
				"calls_p_2025_01",
			},
			expectDrop: []string{
				"public.calls_p_2024_07",
				"public.calls_p_2024_08",
			},
			expectSweep: []timeRange{
				timeRange{start: dateInUTC("2024-07-01"), end: dateInUTC("2024-08-01")},
				timeRange{start: dateInUTC("2024-08-01"), end: dateInUTC("2024-09-01")},
			},
			expectCleanup: []timeRange{
				timeRange{start: dateInUTC("2024-07-01"), end: dateInUTC("2024-08-01")},
				timeRange{start: dateInUTC("2024-08-01"), end: dateInUTC("2024-09-01")},
			},
			expectDetach: []string{
				"public.calls_p_2024_07",
				"public.calls_p_2024_08",
			},
		},
	}

	for _, tc := range tests {
		t.Run(tc.name, func(t *testing.T) {
			db := mocks.NewStore(t)
			dbtx := mocks.NewDBTX(t)
			createdPartitions := make([]string, 0, len(tc.expectCreate))
			sweptRanges := make([]timeRange, 0, len(tc.expectSweep))
			droppedPartitions := make([]string, 0, len(tc.expectDrop))
			cleanupRanges := make([]timeRange, 0, len(tc.expectCleanup))
			detachedPartitions := make([]string, 0, len(tc.expectDetach))

			db.EXPECT().CreatePartition(mctx, mock.AnythingOfType("string"), mock.AnythingOfType("string"), mock.AnythingOfType("time.Time"), mock.AnythingOfType("time.Time")).Run(func(ctx context.Context, tableName, partitionName string, start, end time.Time) {
				createdPartitions = append(createdPartitions, partitionName)
			}).Return(nil)
			db.EXPECT().SweepCalls(mctx, mock.AnythingOfType("pgtype.Timestamptz"), mock.AnythingOfType("pgtype.Timestamptz")).Run(func(ctx context.Context, start, end pgtype.Timestamptz) {
				sweptRanges = append(sweptRanges, timeRange{start: start.Time, end: end.Time})
			}).Return(nil)
			db.EXPECT().CleanupSweptCalls(mctx, mock.AnythingOfType("pgtype.Timestamptz"), mock.AnythingOfType("pgtype.Timestamptz")).Run(func(ctx context.Context, start, end pgtype.Timestamptz) {
				cleanupRanges = append(cleanupRanges, timeRange{start: start.Time, end: end.Time})
			}).Return(nil)
			db.EXPECT().DropPartition(mctx, mock.AnythingOfType("string")).Run(func(ctx context.Context, partName string) {
				droppedPartitions = append(droppedPartitions, partName)
			}).Return(nil)
			db.EXPECT().DetachPartition(mctx, mock.AnythingOfType("string")).Run(func(ctx context.Context, partName string) {
				detachedPartitions = append(detachedPartitions, partName)
			}).Return(nil)
			inTx(db, dbtx)
			db.EXPECT().GetTablePartitions(mctx, "public", "calls").Return(tc.extant, nil)
			pm, err := partman.New(db, tc.cfg)
			require.NoError(t, err)

			err = pm.Check(ctx, tc.now)
			require.NoError(t, err)

			assert.ElementsMatch(t, tc.expectCreate, createdPartitions)
			assert.ElementsMatch(t, tc.expectSweep, sweptRanges)
			assert.ElementsMatch(t, tc.expectDrop, droppedPartitions)
			assert.ElementsMatch(t, tc.expectCleanup, cleanupRanges)
			assert.ElementsMatch(t, tc.expectDetach, detachedPartitions)
		})
	}
}