mantle/internal/budget/budget_test.go at 096bdcebca45016b5e1661b7f38d1332c8e269f8 · dvflw/mantle · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
package budget_test

import (
	"context"
	"testing"
	"time"

	"github.com/dvflw/mantle/internal/budget"
	"github.com/stretchr/testify/assert"
)

func TestCurrentPeriodStart_Calendar(t *testing.T) {
	now := time.Date(2026, 3, 23, 14, 30, 0, 0, time.UTC)
	start := budget.CurrentPeriodStart(now, "calendar", 1)
	assert.Equal(t, time.Date(2026, 3, 1, 0, 0, 0, 0, time.UTC), start)
}

func TestCurrentPeriodStart_Rolling(t *testing.T) {
	now := time.Date(2026, 3, 10, 14, 30, 0, 0, time.UTC)
	start := budget.CurrentPeriodStart(now, "rolling", 15)
	assert.Equal(t, time.Date(2026, 2, 15, 0, 0, 0, 0, time.UTC), start)

	now = time.Date(2026, 3, 20, 14, 30, 0, 0, time.UTC)
	start = budget.CurrentPeriodStart(now, "rolling", 15)
	assert.Equal(t, time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC), start)

	now = time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC)
	start = budget.CurrentPeriodStart(now, "rolling", 15)
	assert.Equal(t, time.Date(2026, 3, 15, 0, 0, 0, 0, time.UTC), start)
}

func TestCurrentPeriodStart_RollingDay28_February(t *testing.T) {
	now := time.Date(2026, 3, 10, 0, 0, 0, 0, time.UTC)
	start := budget.CurrentPeriodStart(now, "rolling", 28)
	assert.Equal(t, time.Date(2026, 2, 28, 0, 0, 0, 0, time.UTC), start)
}

func TestChecker_Check_GlobalHardBlock(t *testing.T) {
	checker := &budget.Checker{
		GlobalMonthlyTokenLimit: 1000,
		ResetMode:               "calendar",
		ResetDay:                1,
		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
			return 1001, nil
		},
	}

	result := checker.Check(context.Background(), budget.CheckInput{
		TeamID:   "team-1",
		Provider: "openai",
	})

	assert.True(t, result.Blocked)
	assert.Equal(t, "global", result.BlockedBy)
	assert.Contains(t, result.Message, "global monthly token limit")
}

func TestChecker_Check_TeamWarnOnly(t *testing.T) {
	checker := &budget.Checker{
		ResetMode: "calendar",
		ResetDay:  1,
		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
			return 0, nil
		},
		GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
			return 5001, nil
		},
		GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
			return &budget.TeamBudget{MonthlyTokenLimit: 5000, Enforcement: "warn"}, nil
		},
	}

	result := checker.Check(context.Background(), budget.CheckInput{
		TeamID:   "team-1",
		Provider: "openai",
	})

	assert.False(t, result.Blocked)
	assert.True(t, result.Warning)
	assert.Contains(t, result.Message, "team budget exceeded")
}

func TestChecker_Check_AllPass(t *testing.T) {
	checker := &budget.Checker{
		GlobalMonthlyTokenLimit: 100000,
		ResetMode:               "calendar",
		ResetDay:                1,
		GetTotalUsage: func(ctx context.Context, period time.Time) (int64, error) {
			return 500, nil
		},
		GetProviderUsage: func(ctx context.Context, teamID, provider string, period time.Time) (int64, error) {
			return 200, nil
		},
		GetTeamBudget: func(ctx context.Context, teamID, provider string) (*budget.TeamBudget, error) {
			return &budget.TeamBudget{MonthlyTokenLimit: 10000, Enforcement: "hard"}, nil
		},
	}

	result := checker.Check(context.Background(), budget.CheckInput{
		TeamID:   "team-1",
		Provider: "openai",
	})

	assert.False(t, result.Blocked)
	assert.False(t, result.Warning)
}

func TestChecker_CheckExecutionBudget(t *testing.T) {
	result := budget.CheckExecutionBudget(50000, 50001)
	assert.True(t, result.Blocked)
	assert.Equal(t, "workflow", result.BlockedBy)

	result = budget.CheckExecutionBudget(50000, 49999)
	assert.False(t, result.Blocked)

	result = budget.CheckExecutionBudget(0, 999999)
	assert.False(t, result.Blocked)
}