package main import ( "errors" "path/filepath" "testing" "time" ) type clock struct{ now time.Time } func (c *clock) Now() time.Time { return c.now } func testLimiter(t *testing.T, cfg Config, at string) (*Limiter, *clock) { t.Helper() moment, err := time.Parse(time.RFC3339, at) if err != nil { t.Fatal(err) } cfg.Timezone = "UTC" c := &clock{now: moment} limiter, err := NewLimiter(cfg, "", c.Now) if err != nil { t.Fatal(err) } return limiter, c } func TestInWindowTable(t *testing.T) { at := func(text string) time.Time { moment, _ := time.Parse("15:04", text) return moment } cases := []struct { from, to, now string want bool }{ {"", "", "12:00", true}, {"09:00", "17:00", "09:00", true}, {"09:00", "17:00", "16:59", true}, {"09:00", "17:00", "17:00", false}, {"09:00", "17:00", "08:59", false}, {"22:00", "06:00", "23:30", true}, {"22:00", "06:00", "05:59", true}, {"22:00", "06:00", "06:00", false}, {"22:00", "06:00", "12:00", false}, {"10:00", "10:00", "03:00", true}, } for _, item := range cases { if got := InWindow(item.from, item.to, at(item.now)); got != item.want { t.Errorf("%s-%s at %s: got %v", item.from, item.to, item.now, got) } } } func TestAcquireConcurrencyCapAndServerCannotWiden(t *testing.T) { cfg := DefaultConfig() cfg.MaxConcurrent = 2 limiter, _ := testLimiter(t, cfg, "2026-10-15T12:00:00Z") limiter.ApplyServer(&ServerLimits{MaxConcurrentPerNode: 8}, &ServerSettings{MaxConcurrent: 8}) if limiter.Acquire() != "" || limiter.Acquire() != "" { t.Fatal("two slots should be free") } if got := limiter.Acquire(); got != RefuseConcurrency { t.Fatalf("third acquire: %q", got) } limiter.Release() if got := limiter.Acquire(); got != "" { t.Fatalf("after release: %q", got) } } func TestServerCanNarrowConcurrency(t *testing.T) { cfg := DefaultConfig() cfg.MaxConcurrent = 4 limiter, _ := testLimiter(t, cfg, "2026-10-15T12:00:00Z") limiter.ApplyServer(&ServerLimits{MaxConcurrentPerNode: 2}, &ServerSettings{MaxConcurrent: 1}) if limiter.Acquire() != "" { t.Fatal("first slot") } if got := limiter.Acquire(); got != RefuseConcurrency { t.Fatalf("got %q", got) } } func TestBandwidthCapLocalAndDayRollover(t *testing.T) { cfg := DefaultConfig() cfg.BandwidthMBDay = 1 limiter, c := testLimiter(t, cfg, "2026-10-15T23:00:00Z") limiter.ApplyServer(&ServerLimits{MaxDonorMBPerDay: 20480}, &ServerSettings{BandwidthMBDay: 9999}) if limiter.Acquire() != "" { t.Fatal("fresh day must allow") } limiter.Release() limiter.AddBytes(mebibyte - 10) if got := limiter.Remaining(); got != 10 { t.Fatalf("remaining %d", got) } limiter.AddBytes(10) if got := limiter.Acquire(); got != RefuseBandwidth { t.Fatalf("got %q", got) } c.now = c.now.Add(2 * time.Hour) if got := limiter.Acquire(); got != "" { t.Fatalf("next day: %q", got) } } func TestServerCanNarrowBandwidth(t *testing.T) { cfg := DefaultConfig() cfg.BandwidthMBDay = 500 limiter, _ := testLimiter(t, cfg, "2026-10-15T12:00:00Z") limiter.ApplyServer(nil, &ServerSettings{BandwidthMBDay: 2}) if got := limiter.Remaining(); got != 2*mebibyte { t.Fatalf("remaining %d", got) } } func TestActiveHoursLocalAndServer(t *testing.T) { cfg := DefaultConfig() cfg.ActiveFrom, cfg.ActiveTo = "22:00", "06:00" limiter, c := testLimiter(t, cfg, "2026-10-15T12:00:00Z") if got := limiter.Acquire(); got != RefuseHours { t.Fatalf("noon: %q", got) } c.now = c.now.Add(11 * time.Hour) if got := limiter.Acquire(); got != "" { t.Fatalf("23:00: %q", got) } limiter.Release() limiter.ApplyServer(nil, &ServerSettings{ActiveFrom: "01:00", ActiveTo: "02:00", Timezone: "UTC"}) if got := limiter.Acquire(); got != RefuseHours { t.Fatalf("server window: %q", got) } } func TestUsagePersistsAcrossRestart(t *testing.T) { path := filepath.Join(t.TempDir(), "usage.json") cfg := DefaultConfig() cfg.Timezone = "UTC" moment := time.Date(2026, 10, 15, 12, 0, 0, 0, time.UTC) c := &clock{now: moment} first, err := NewLimiter(cfg, path, c.Now) if err != nil { t.Fatal(err) } first.AddBytes(12345) first.Flush() second, err := NewLimiter(cfg, path, c.Now) if err != nil { t.Fatal(err) } if _, used, _ := second.Usage(); used != 12345 { t.Fatalf("used %d", used) } c.now = moment.Add(24 * time.Hour) if _, used, _ := second.Usage(); used != 0 { t.Fatalf("next day used %d", used) } if got := ReadUsage(path, time.UTC, moment); got != 12345 { t.Fatalf("read usage %d", got) } } func TestAcquireRefusesOnBatteryOnlyWhenAsked(t *testing.T) { limiter, c := testLimiter(t, DefaultConfig(), "2026-01-15T12:00:00Z") calls := 0 onBattery := true limiter.SetBatteryProbe(func() (bool, error) { calls++; return onBattery, nil }, nil) if got := limiter.Acquire(); got != "" { t.Fatalf("setting off must ignore the probe: %q", got) } limiter.Release() if calls != 0 { t.Fatalf("probe ran %d times with the setting off", calls) } limiter.ApplyServer(nil, &ServerSettings{PauseOnBattery: true}) if got := limiter.Acquire(); got != RefuseBattery { t.Fatalf("got %q", got) } onBattery = false if got := limiter.Acquire(); got != RefuseBattery || calls != 1 { t.Fatalf("cached result expected: %q calls %d", got, calls) } c.now = c.now.Add(16 * time.Second) if got := limiter.Acquire(); got != "" || calls != 2 { t.Fatalf("after cache expiry: %q calls %d", got, calls) } } func TestBatteryProbeErrorMeansMainsAndLogsOnce(t *testing.T) { limiter, c := testLimiter(t, DefaultConfig(), "2026-01-15T12:00:00Z") logged := 0 limiter.SetBatteryProbe(func() (bool, error) { return true, errors.New("no sysfs") }, func(string, ...any) { logged++ }) limiter.ApplyServer(nil, &ServerSettings{PauseOnBattery: true}) for i := 0; i < 3; i++ { if got := limiter.Acquire(); got != "" { t.Fatalf("probe error must not refuse: %q", got) } limiter.Release() c.now = c.now.Add(20 * time.Second) } if logged != 1 { t.Fatalf("logged %d times", logged) } } func TestBatteryHelpers(t *testing.T) { if !parsePmset("Now drawing from 'Battery Power'\n -InternalBattery-0") || parsePmset("Now drawing from 'AC Power'") { t.Fatal("pmset parse") } cases := []struct { ac, flag byte want bool }{{0, 1, true}, {0, 8, true}, {1, 1, false}, {0, 128, false}, {255, 255, false}, {0, 255, false}} for _, item := range cases { if got := windowsOnBattery(item.ac, item.flag); got != item.want { t.Errorf("ac %d flag %d: %v", item.ac, item.flag, got) } } } const ( parityBase = 1767225600 parityStep = 40009 paritySamples = 200 ) var parityCases = []struct{ from, to, zone, bits string }{ {"09:00", "17:00", "Europe/Amsterdam", "01010000001010101010000101010101000000101010100000010101010000001010101010000101010101000000101010100000010101010000001010101010000101010101000000101010100000010101010000001010101010000001010101010000"}, {"22:00", "06:00", "America/New_York", "00010101010000001010101000000101010101000010101010100000010101010000001010101000000101010101000010101010100000010101010000001010101000000101010101000000101010101000010101010100000010101010000001010101"}, {"08:30", "12:15", "Asia/Kolkata", "00000001010000000000101000000000010101000000000010100000000001010000000000101000000000000101000000000010100000000001010000000000101000000000000101000000000010100000000001010000000000101010000000000101"}, {"00:00", "12:00", "Pacific/Auckland", "01101010101010110101010101010010101010101001010101010101101010101010110101010101010010101010101001010101010101101010101010110101010101010010101010101011010101010101101010101010100101010101010010101010"}, {"10:00", "10:00", "UTC", "11111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111"}, {"09:00", "17:00", "Not/AZone", "01010000101010101000000101010100000010101010000001010101010000001010101000000101010100000010101010100001010101010000001010101000000101010100000010101010100001010101010000001010101000000101010100000010"}, {"23:00", "01:00", "Australia/Lord_Howe", "00000000000010000000000001010000000000001000000000000100000000000010000000000001010000000000001000000000000100000000000010000000000000010000000000001000000000000100000000000010000000000000010000000000"}, {"", "", "Europe/Amsterdam", "11111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111111"}, } func TestInWindowParityWithPythonScheduler(t *testing.T) { for _, item := range parityCases { zone, err := time.LoadLocation(item.zone) if err != nil { zone = time.UTC } for index := 0; index < paritySamples; index++ { moment := time.Unix(int64(parityBase+index*parityStep), 0).In(zone) if got := InWindow(item.from, item.to, moment); got != (item.bits[index] == '1') { t.Fatalf("%s-%s %s #%d: got %v", item.from, item.to, item.zone, index, got) } } } }