package budget import ( "errors" "io" "path/filepath" "reflect" "strings" "testing" "github.com/godbus/dbus/v5" ) func TestSearchScope(t *testing.T) { props, err := scopeProps(testConfig(), "search", "", []uint32{10, 11}) if err != nil { t.Fatal(err) } got := values(props) if got["Slice"].Value() != "search.slice" { t.Fatal(got["Slice"]) } if _, exists := got["MemoryMax"]; exists { t.Fatal("search scopes should share the slice ceiling") } if got["PIDs"].Signature().String() != "au" || !reflect.DeepEqual(got["PIDs"].Value(), []uint32{10, 11}) { t.Fatal(got["PIDs"]) } props, err = scopeProps(testConfig(), "default", "app-editor.slice", []uint32{10}) if err != nil { t.Fatal(err) } got = values(props) if got["Slice"].Value() != "app-editor.slice" || got["MemoryMax"].Value() != uint64(8589934592) { t.Fatal(got) } } func TestWatcherFamily(t *testing.T) { group := "/app.slice/app-terminal-123.scope" p := fakeProcs{map[uint32]task{10: {1, group, "compute", ""}, 11: {10, group, "search", ""}}, map[uint32]string{10: group, 11: group}} f := newFake() f.onStart = func(name string, props []property) { for _, pid := range values(props)["PIDs"].Value().([]uint32) { p.groups[pid] = "/app.slice/" + name } } w := watcher{config: testConfig(), ctl: f, procs: p, cgroup: t.TempDir(), out: io.Discard} if err := w.init(); err != nil { t.Fatal(err) } if err := w.step(); err != nil { t.Fatal(err) } if len(f.starts) != 1 { t.Fatalf("family was assigned twice: %v", f.starts) } for name, props := range f.starts { if len(name) > 20 || !strings.HasPrefix(name, "budget-c") { t.Fatal(name) } if !reflect.DeepEqual(values(props)["PIDs"].Value(), []uint32{10, 11}) { t.Fatal(props) } } p.tasks = map[uint32]task{10: {1, p.groups[10], "", ""}, 11: {10, p.groups[11], "search", ""}} w.procs = p if err := w.step(); err != nil { t.Fatal(err) } if len(f.starts) != 2 { t.Fatal(f.starts) } for name, props := range f.starts { if strings.HasPrefix(name, "budget-s") && values(props)["Slice"].Value() != "search.slice" { t.Fatal(props) } } } func TestWatcherRaces(t *testing.T) { group := "/app.slice/app-terminal-123.scope" for _, tc := range []struct { name, current string err error }{ {"moved", "/other", nil}, {"exited", "", nil}, {"API exit", group, dbus.Error{Name: noPID}}, } { t.Run(tc.name, func(t *testing.T) { p := fakeProcs{map[uint32]task{10: {1, group, "search", ""}}, map[uint32]string{}} if tc.current != "" { p.groups[10] = tc.current } f := newFake() f.startErr = tc.err w := watcher{config: testConfig(), ctl: f, procs: p, cgroup: t.TempDir(), out: io.Discard} if err := w.init(); err != nil { t.Fatal(err) } if err := w.step(); err != nil { t.Fatal(err) } if len(w.owned) != 0 || len(f.starts) != 0 { t.Fatal("an exited or moved task was retained") } }) } f := newFake() f.startErr = errors.New("bus unavailable") w := watcher{config: testConfig(), ctl: f, procs: fakeProcs{map[uint32]task{10: {1, group, "search", ""}}, map[uint32]string{10: group}}, cgroup: t.TempDir(), owned: map[string]bool{}, out: io.Discard} if err := w.step(); err != f.startErr { t.Fatalf("unexpected error was hidden: %v", err) } } func TestReap(t *testing.T) { f := newFake() f.states = map[string]unitState{ "budget-c000001.scope": {"active", "/empty"}, "budget-c000002.scope": {"active", "/live"}, "budget-c000003.scope": {"activating", ""}, "budget-c000004.scope": {"inactive", ""}, "budget-c000005.scope": {"failed", ""}, "budget-c000007.scope": {"active", "/removed"}, } root := t.TempDir() put(t, filepath.Join(root, "empty/cgroup.events"), "populated 0\nfrozen 0\n") put(t, filepath.Join(root, "live/cgroup.events"), "populated 1\nfrozen 0\n") w := watcher{ctl: f, cgroup: root, owned: map[string]bool{}} for i := 1; i <= 7; i++ { w.owned["budget-c00000"+string(rune('0'+i))+".scope"] = true } if err := w.reap(); err != nil { t.Fatal(err) } if !reflect.DeepEqual(f.stops, []string{"budget-c000001.scope"}) { t.Fatal(f.stops) } if !reflect.DeepEqual(keys(w.owned), []string{"budget-c000002.scope", "budget-c000003.scope"}) { t.Fatal(w.owned) } f.stateErr["budget-c000002.scope"] = errors.New("bus unavailable") if err := w.reap(); err == nil { t.Fatal("unexpected bus error hidden") } } func TestRecoverScopes(t *testing.T) { f := newFake() f.units = []unit{{Name: "budget-c123abc.scope"}, {Name: "budget.service"}, {Name: "work-123.scope"}, {Name: "app-terminal-123.scope"}} w := watcher{ctl: f} if err := w.init(); err != nil { t.Fatal(err) } if !reflect.DeepEqual(keys(w.owned), []string{"budget-c123abc.scope"}) { t.Fatal(w.owned) } } func TestWatcherOutput(t *testing.T) { want := errors.New("output failed") group := "/app.slice/app-terminal-123.scope" w := watcher{ config: testConfig(), ctl: newFake(), cgroup: t.TempDir(), out: failWriter{want}, procs: fakeProcs{map[uint32]task{10: {1, group, "search", ""}}, map[uint32]string{10: group}}, } if err := w.init(); err != nil { t.Fatal(err) } if err := w.step(); !errors.Is(err, want) { t.Fatal(err) } }