tokens-reef/backend/internal/service/ops_partition_test.go

package service

import (
	"context"
	"errors"
	"testing"
	"time"
)

// mockOpsRepoForPartition is a mock for testing partition status
type mockOpsRepoForPartition struct {
	isPartitioned      bool
	isPartitionedErr   error
	rowCount           int64
	rowCountErr        error
	partitionCount     int
	partitionCountErr  error
}

func (m *mockOpsRepoForPartition) IsUsageLogsPartitioned(ctx context.Context) (bool, error) {
	return m.isPartitioned, m.isPartitionedErr
}

func (m *mockOpsRepoForPartition) GetUsageLogsRowCount(ctx context.Context) (int64, error) {
	return m.rowCount, m.rowCountErr
}

func (m *mockOpsRepoForPartition) GetUsageLogsPartitionCount(ctx context.Context) (int, error) {
	return m.partitionCount, m.partitionCountErr
}

// Implement minimal OpsRepository interface for testing
func (m *mockOpsRepoForPartition) InsertErrorLog(ctx context.Context, input *OpsInsertErrorLogInput) (int64, error) {
	return 0, nil
}
func (m *mockOpsRepoForPartition) BatchInsertErrorLogs(ctx context.Context, inputs []*OpsInsertErrorLogInput) (int64, error) {
	return 0, nil
}
func (m *mockOpsRepoForPartition) ListErrorLogs(ctx context.Context, filter *OpsErrorLogFilter) (*OpsErrorLogList, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetErrorLogByID(ctx context.Context, id int64) (*OpsErrorLogDetail, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) ListRequestDetails(ctx context.Context, filter *OpsRequestDetailFilter) ([]*OpsRequestDetail, int64, error) {
	return nil, 0, nil
}
func (m *mockOpsRepoForPartition) BatchInsertSystemLogs(ctx context.Context, inputs []*OpsInsertSystemLogInput) (int64, error) {
	return 0, nil
}
func (m *mockOpsRepoForPartition) ListSystemLogs(ctx context.Context, filter *OpsSystemLogFilter) (*OpsSystemLogList, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) DeleteSystemLogs(ctx context.Context, filter *OpsSystemLogCleanupFilter) (int64, error) {
	return 0, nil
}
func (m *mockOpsRepoForPartition) InsertSystemLogCleanupAudit(ctx context.Context, input *OpsSystemLogCleanupAudit) error {
	return nil
}
func (m *mockOpsRepoForPartition) InsertRetryAttempt(ctx context.Context, input *OpsInsertRetryAttemptInput) (int64, error) {
	return 0, nil
}
func (m *mockOpsRepoForPartition) UpdateRetryAttempt(ctx context.Context, input *OpsUpdateRetryAttemptInput) error {
	return nil
}
func (m *mockOpsRepoForPartition) GetLatestRetryAttemptForError(ctx context.Context, sourceErrorID int64) (*OpsRetryAttempt, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) ListRetryAttemptsByErrorID(ctx context.Context, sourceErrorID int64, limit int) ([]*OpsRetryAttempt, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) UpdateErrorResolution(ctx context.Context, errorID int64, resolved bool, resolvedByUserID *int64, resolvedRetryID *int64, resolvedAt *time.Time) error {
	return nil
}
func (m *mockOpsRepoForPartition) GetWindowStats(ctx context.Context, filter *OpsDashboardFilter) (*OpsWindowStats, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetRealtimeTrafficSummary(ctx context.Context, filter *OpsDashboardFilter) (*OpsRealtimeTrafficSummary, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetDashboardOverview(ctx context.Context, filter *OpsDashboardFilter) (*OpsDashboardOverview, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetThroughputTrend(ctx context.Context, filter *OpsDashboardFilter, bucketSeconds int) (*OpsThroughputTrendResponse, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetLatencyHistogram(ctx context.Context, filter *OpsDashboardFilter) (*OpsLatencyHistogramResponse, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetErrorTrend(ctx context.Context, filter *OpsDashboardFilter, bucketSeconds int) (*OpsErrorTrendResponse, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetErrorDistribution(ctx context.Context, filter *OpsDashboardFilter) (*OpsErrorDistributionResponse, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetOpenAITokenStats(ctx context.Context, filter *OpsOpenAITokenStatsFilter) (*OpsOpenAITokenStatsResponse, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) InsertSystemMetrics(ctx context.Context, input *OpsInsertSystemMetricsInput) error {
	return nil
}
func (m *mockOpsRepoForPartition) GetLatestSystemMetrics(ctx context.Context, windowMinutes int) (*OpsSystemMetricsSnapshot, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) UpsertJobHeartbeat(ctx context.Context, input *OpsUpsertJobHeartbeatInput) error {
	return nil
}
func (m *mockOpsRepoForPartition) ListJobHeartbeats(ctx context.Context) ([]*OpsJobHeartbeat, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) ListAlertRules(ctx context.Context) ([]*OpsAlertRule, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) CreateAlertRule(ctx context.Context, input *OpsAlertRule) (*OpsAlertRule, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) UpdateAlertRule(ctx context.Context, input *OpsAlertRule) (*OpsAlertRule, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) DeleteAlertRule(ctx context.Context, id int64) error {
	return nil
}
func (m *mockOpsRepoForPartition) ListAlertEvents(ctx context.Context, filter *OpsAlertEventFilter) ([]*OpsAlertEvent, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetAlertEventByID(ctx context.Context, eventID int64) (*OpsAlertEvent, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetActiveAlertEvent(ctx context.Context, ruleID int64) (*OpsAlertEvent, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) GetLatestAlertEvent(ctx context.Context, ruleID int64) (*OpsAlertEvent, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) CreateAlertEvent(ctx context.Context, event *OpsAlertEvent) (*OpsAlertEvent, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) UpdateAlertEventStatus(ctx context.Context, eventID int64, status string, resolvedAt *time.Time) error {
	return nil
}
func (m *mockOpsRepoForPartition) UpdateAlertEventEmailSent(ctx context.Context, eventID int64, emailSent bool) error {
	return nil
}
func (m *mockOpsRepoForPartition) CreateAlertSilence(ctx context.Context, input *OpsAlertSilence) (*OpsAlertSilence, error) {
	return nil, nil
}
func (m *mockOpsRepoForPartition) IsAlertSilenced(ctx context.Context, ruleID int64, platform string, groupID *int64, region *string, now time.Time) (bool, error) {
	return false, nil
}
func (m *mockOpsRepoForPartition) UpsertHourlyMetrics(ctx context.Context, startTime, endTime time.Time) error {
	return nil
}
func (m *mockOpsRepoForPartition) UpsertDailyMetrics(ctx context.Context, startTime, endTime time.Time) error {
	return nil
}
func (m *mockOpsRepoForPartition) GetLatestHourlyBucketStart(ctx context.Context) (time.Time, bool, error) {
	return time.Time{}, false, nil
}
func (m *mockOpsRepoForPartition) GetLatestDailyBucketDate(ctx context.Context) (time.Time, bool, error) {
	return time.Time{}, false, nil
}

func TestGetUsageLogsPartitionStatus_Partitioned(t *testing.T) {
	mock := &mockOpsRepoForPartition{
		isPartitioned:  true,
		rowCount:       500000,
		partitionCount: 12,
	}

	svc := &OpsService{opsRepo: mock}

	status, err := svc.GetUsageLogsPartitionStatus(context.Background())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if !status.IsPartitioned {
		t.Error("expected IsPartitioned to be true")
	}
	if status.RowCount != 500000 {
		t.Errorf("expected RowCount 500000, got %d", status.RowCount)
	}
	if status.PartitionCount != 12 {
		t.Errorf("expected PartitionCount 12, got %d", status.PartitionCount)
	}
	if status.WarningLevel != "none" {
		t.Errorf("expected WarningLevel 'none', got %s", status.WarningLevel)
	}
	if status.NeedsPartitioning {
		t.Error("expected NeedsPartitioning to be false")
	}
}

func TestGetUsageLogsPartitionStatus_NotPartitioned_NeedsPartitioning(t *testing.T) {
	mock := &mockOpsRepoForPartition{
		isPartitioned: false,
		rowCount:      150000, // Above threshold of 100000
	}

	svc := &OpsService{opsRepo: mock}

	status, err := svc.GetUsageLogsPartitionStatus(context.Background())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if status.IsPartitioned {
		t.Error("expected IsPartitioned to be false")
	}
	if !status.NeedsPartitioning {
		t.Error("expected NeedsPartitioning to be true")
	}
	if status.WarningLevel != "warning" {
		t.Errorf("expected WarningLevel 'warning', got %s", status.WarningLevel)
	}
}

func TestGetUsageLogsPartitionStatus_NotPartitioned_BelowThreshold(t *testing.T) {
	mock := &mockOpsRepoForPartition{
		isPartitioned: false,
		rowCount:      30000, // Well below threshold (50000)
	}

	svc := &OpsService{opsRepo: mock}

	status, err := svc.GetUsageLogsPartitionStatus(context.Background())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if status.NeedsPartitioning {
		t.Error("expected NeedsPartitioning to be false")
	}
	if status.WarningLevel != "none" {
		t.Errorf("expected WarningLevel 'none', got %s", status.WarningLevel)
	}
}

func TestGetUsageLogsPartitionStatus_NotPartitioned_InfoLevel(t *testing.T) {
	mock := &mockOpsRepoForPartition{
		isPartitioned: false,
		rowCount:      75000, // Between 50000 and 100000
	}

	svc := &OpsService{opsRepo: mock}

	status, err := svc.GetUsageLogsPartitionStatus(context.Background())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if status.NeedsPartitioning {
		t.Error("expected NeedsPartitioning to be false")
	}
	if status.WarningLevel != "info" {
		t.Errorf("expected WarningLevel 'info', got %s", status.WarningLevel)
	}
}

func TestGetUsageLogsPartitionStatus_Error(t *testing.T) {
	mock := &mockOpsRepoForPartition{
		isPartitionedErr: errors.New("db error"),
	}

	svc := &OpsService{opsRepo: mock}

	_, err := svc.GetUsageLogsPartitionStatus(context.Background())
	if err == nil {
		t.Error("expected error, got nil")
	}
}

func TestGetUsageLogsPartitionStatus_NilRepo(t *testing.T) {
	svc := &OpsService{opsRepo: nil}

	_, err := svc.GetUsageLogsPartitionStatus(context.Background())
	if err == nil {
		t.Error("expected error for nil repo, got nil")
	}
}