chultun/internal/budget/watch_test.go

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)
	}
}