chultun/internal/budget/units_test.go

package budget

import (
	"strings"
	"testing"
)

func TestEscape(t *testing.T) {
	for input, want := range map[string]string{
		"editor-uri-handler": `editor\x2duri\x2dhandler`,
		".hidden":            `\x2ehidden`, "a/b c": `a-b\x20c`, "λ": `\xce\xbb`,
		"org.example.Terminal": "org.example.Terminal", "a:b_c": "a:b_c",
	} {
		if got := escape(input); got != want {
			t.Errorf("%q: got %q, want %q", input, got, want)
		}
	}
}

func TestUnitProfiles(t *testing.T) {
	c := testConfig()
	for name, want := range map[string]string{
		"app-terminal-surface-transient-123.scope": "unrestricted",
		"app-org.example.Terminal.service":         "unrestricted",
		"[email protected]":      "unrestricted",
		"containerd.service":                       "unrestricted", "worker.service": "unrestricted",
		"[email protected]":                                   "compute",
		"app-flatpak-org.example.Game-123.scope":                 "compute",
		"app-flatpak-org.example.Download-123.scope":             "bulk",
		`app-flatpak-org.example.Browser\x2dcommunity-123.scope`: "default",
		"[email protected]":                                 "default", "app-editor.slice": "default",
		"search.slice": "search", "session-ui.service": "",
		"budget.service": "", "session.slice": "", "work-123.scope": "",
	} {
		t.Run(name, func(t *testing.T) {
			if got := c.unitProfile(name); got != want {
				t.Fatalf("got %q, want %q", got, want)
			}
		})
	}
}

func TestRender(t *testing.T) {
	c := testConfig()
	files, err := c.render()
	if err != nil {
		t.Fatal(err)
	}
	for _, name := range []string{"app-.scope.d/50-budget.conf", "app-.service.d/50-budget.conf"} {
		if !strings.Contains(files[name], "MemoryMax=8G\n") {
			t.Error(name)
		}
	}
	for _, app := range []string{"editor", "editor-uri-handler"} {
		for _, directory := range appDirs(app) {
			if !strings.Contains(files[directory+"/50-budget.conf"], "Slice=app-editor.slice\n") {
				t.Error(directory)
			}
		}
	}
	if !strings.Contains(files["search.slice"], "CPUQuota=200%\n") || !strings.Contains(files["search.slice"], "IOReadBandwidthMax=/dev/mock 50M\n") {
		t.Fatal(files["search.slice"])
	}
	if !strings.Contains(files["app-editor.slice"], "MemoryMax=8G\n") {
		t.Fatal(files["app-editor.slice"])
	}
	if !strings.Contains(files["app-terminal-.scope.d/50-budget.conf"], "CPUQuota=\n") {
		t.Fatal("terminal ceiling wasn't removed")
	}
	if !strings.Contains(files["app-flatpak-org.example.Game-.scope.d/50-budget.conf"], "MemoryMax=16G\n") {
		t.Fatal("compute profile missing")
	}
	if files["budget.service"] != serviceText {
		t.Fatal("wrong service command")
	}
	c.Shared["editor"] = "missing"
	if _, err := c.render(); err == nil {
		t.Fatal("unknown shared profile accepted")
	}
}

func TestNoSharedApps(t *testing.T) {
	c := testConfig()
	c.Shared = nil
	if got := c.sharedSlice(); got != "" {
		t.Fatal(got)
	}
	files, err := c.render()
	if err != nil {
		t.Fatal(err)
	}
	for name := range files {
		if strings.HasPrefix(name, "app-editor") {
			t.Fatal("unconfigured shared slice remains", name)
		}
	}
	if got := c.unitProfile(""); got != "" {
		t.Fatal(got)
	}
}

func TestSliceDropIns(t *testing.T) {
	c := testConfig()
	c.Units["compute"] = []string{"custom.slice"}
	c.Prefixes["bulk"] = []string{"custom-.slice.d"}
	files, err := c.render()
	if err != nil {
		t.Fatal(err)
	}
	for file, limit := range map[string]string{
		"custom.slice.d/50-budget.conf": "16G", "custom-.slice.d/50-budget.conf": "8G",
	} {
		if text := files[file]; !strings.HasPrefix(text, "[Slice]\n") || !strings.Contains(text, "MemoryMax="+limit+"\n") {
			t.Fatalf("%s: %q", file, text)
		}
	}
	if got := c.unitProfile("custom-work.slice"); got != "bulk" {
		t.Fatalf("slice prefix profile: %q", got)
	}
}

func TestInvalidPrefixes(t *testing.T) {
	for _, prefix := range []string{"snap-firefox", "snap.firefox", "snap.firefox.scope", "snap.firefox.timer.d"} {
		t.Run(prefix, func(t *testing.T) {
			c := testConfig()
			c.Prefixes["default"] = []string{prefix}
			if _, err := c.render(); err == nil {
				t.Fatal("invalid prefix accepted")
			}
			if got := c.unitProfile("session-ui.service"); got != "" {
				t.Fatalf("invalid prefix matched an unrelated unit: %q", got)
			}
		})
	}
}