package budget import ( "errors" "fmt" "io" "os" "path/filepath" "reflect" "strings" "testing" ) func TestApply(t *testing.T) { f := newFake() f.units = []unit{ {Name: "app-browser-123.scope", Active: "active"}, {Name: "app-org.example.Terminal.service", Active: "active"}, {Name: "containerd.service", Active: "active"}, {Name: "session-ui.service", Active: "active"}, {Name: "app-closed-123.scope", Active: "inactive"}, } events := []string{} f.events = &events a := application{testConfig(), homePaths(t.TempDir()), f, func(args ...string) error { events = append(events, strings.Join(args, " ")) return nil }, io.Discard} if err := a.apply(); err != nil { t.Fatal(err) } want := []string{"daemon-reload", "set app-browser-123.scope", "set app-org.example.Terminal.service", "set containerd.service", "enable budget.service", "restart budget.service"} if !reflect.DeepEqual(events, want) { t.Fatalf("got %v, want %v", events, want) } if values(f.sets["app-browser-123.scope"])["MemoryMax"].Value() != uint64(8589934592) { t.Fatal("default ceiling missing") } if len(f.sets) != 3 { t.Fatal("session infrastructure or inactive apps were changed") } f.setErr = errors.New("set failed") events = nil if err := a.apply(); err != f.setErr { t.Fatal(err) } if len(events) != 2 { t.Fatal("continued after failure", events) } } func TestApplyShared(t *testing.T) { c := testConfig() c.Shared["editor"] = "compute" f := newFake() f.units = []unit{{Name: "app-editor@123.service", Active: "active"}} a := application{c, homePaths(t.TempDir()), f, func(...string) error { return nil }, io.Discard} if err := a.apply(); err != nil { t.Fatal(err) } if got := values(f.sets[f.units[0].Name])["MemoryMax"].Value(); got != uint64(17179869184) { t.Fatalf("shared app ceiling: got %v, want 16 GiB", got) } data, err := os.ReadFile(filepath.Join(a.paths.units, "app-editor@.service.d/50-budget.conf")) if err != nil || !strings.Contains(string(data), "MemoryMax=16G\n") { t.Fatalf("persistent shared ceiling: %q, %v", data, err) } } func TestApplyScopes(t *testing.T) { c := testConfig() c.Profiles["compute"] = profile{"400%", "2G", "4G", "1G"} c.Profiles["bulk"] = profile{"200%", "1G", "2G", "1G"} c.Profiles["cache"] = profile{"100%", "1G", "3G", "0"} f := newFake() f.units = []unit{ {Name: "budget-c123abc.scope", Description: "compute workload", Active: "active"}, {Name: "budget-b123abc.scope", Description: "bulk workload", Active: "active"}, {Name: "budget-c456abc.scope", Description: "cache workload", Active: "active"}, {Name: "budget-s123abc.scope", Description: "search workload", Active: "active"}, {Name: "budget-c123def.scope", Description: "compute workload", Active: "inactive"}, } a := application{c, homePaths(t.TempDir()), f, func(...string) error { return nil }, io.Discard} if err := a.apply(); err != nil { t.Fatal(err) } for name, want := range map[string]uint64{ "budget-c123abc.scope": 4294967296, "budget-b123abc.scope": 2147483648, "budget-c456abc.scope": 3221225472, } { if got := values(f.sets[name])["MemoryMax"].Value(); got != want { t.Errorf("%s ceiling: got %v, want %d", name, got, want) } } if len(f.sets) != 3 { t.Fatalf("unexpected scope updates: %v", keys(f.sets)) } } func TestApplyInvalid(t *testing.T) { for name, change := range map[string]func(*config){ "resource": func(c *config) { p := c.Profiles["default"] p.High = "invalid" c.Profiles["default"] = p }, "command": func(c *config) { c.Commands["java"] = "missing" }, "prefix": func(c *config) { c.Prefixes["default"] = []string{"snap-firefox"} }, } { for _, active := range []bool{false, true} { t.Run(fmt.Sprintf("%s/active=%t", name, active), func(t *testing.T) { c := testConfig() change(&c) f := newFake() if active { f.units = []unit{{Name: "app-browser-123.scope", Active: "active"}} } var events []string f.events = &events p := homePaths(t.TempDir()) put(t, filepath.Join(p.units, "user.conf"), "keep") a := application{c, p, f, func(args ...string) error { events = append(events, strings.Join(args, " ")) return nil }, io.Discard} if err := a.apply(); err == nil { t.Fatal("invalid configuration accepted") } if len(events) != 0 { t.Fatal("invalid configuration changed units:", events) } entries, err := os.ReadDir(p.units) if err != nil || len(entries) != 1 || entries[0].Name() != "user.conf" { t.Fatalf("invalid configuration changed files: %v, %v", entries, err) } if _, err := os.Stat(p.state); !os.IsNotExist(err) { t.Fatalf("invalid configuration created state: %v", err) } }) } } } func TestSharedOverride(t *testing.T) { f := newFake() f.units = []unit{ {Name: "app-editor@123.service", Active: "active"}, {Name: `app-editor\x2duri\x2dhandler@456.service`, Active: "active"}, {Name: "budget-d123abc.scope", Active: "active"}, {Name: "app-other-123.scope", Active: "active"}, {Name: "app-editor@closed.service", Active: "inactive"}, } f.groups = map[string]string{ "app-editor.slice": "/app.slice/app-editor.slice", "app-editor@123.service": "/app.slice/app-editor.slice/app-editor@123.service", `app-editor\x2duri\x2dhandler@456.service`: "/app.slice/app-editor.slice/uri.service", "budget-d123abc.scope": "/app.slice/app-editor.slice/budget-d123abc.scope", "app-other-123.scope": "/app.slice/app-editor.slice-neighbor/app-other-123.scope", } a := application{config: testConfig(), ctl: f, out: io.Discard} if err := a.override("compute", "editor"); err != nil { t.Fatal(err) } want := []string{"app-editor.slice", "app-editor@123.service", `app-editor\x2duri\x2dhandler@456.service`, "budget-d123abc.scope"} if !reflect.DeepEqual(keys(f.sets), want) { t.Fatalf("got %v, want %v", keys(f.sets), want) } for name, props := range f.sets { if values(props)["MemoryMax"].Value() != uint64(17179869184) { t.Fatal(name, props) } } if err := a.override("default", "app-editor.slice"); err != nil { t.Fatal(err) } for name, props := range f.sets { if values(props)["MemoryMax"].Value() != uint64(8589934592) { t.Fatal(name, props) } } } func TestOverrideTargets(t *testing.T) { f := newFake() f.units = []unit{{Name: "app-flatpak-org.example.Download-123.scope", Active: "active"}} a := application{config: testConfig(), ctl: f, out: io.Discard} if err := a.override("unrestricted", "org.example.Download"); err != nil { t.Fatal(err) } if len(f.sets) != 1 { t.Fatal(f.sets) } if err := a.override("default", "no-such-app"); err == nil { t.Fatal("missing app accepted") } if err := a.override("missing", "org.example.Download"); err == nil { t.Fatal("missing profile accepted") } } func TestAppOutput(t *testing.T) { want := errors.New("output failed") f := newFake() a := application{ config: testConfig(), paths: homePaths(t.TempDir()), ctl: f, run: func(...string) error { return nil }, out: failWriter{want}, } if err := a.apply(); !errors.Is(err, want) { t.Fatal(err) } if err := a.override("default", "app-example.service"); !errors.Is(err, want) { t.Fatal(err) } }