internal/testsuite: additional polling helpers

These are mostly for process lifecycle checks. This change also adds
polling for the Xwayland server in the sway helper.

Signed-off-by: Ophestra <cat@gensokyo.uk>

1 files changed, 97 insertions, 0 deletions

diff --git a/internal/testsuite/testsuite.go b/internal/testsuite/testsuite.go
index 7c198bba..ff6c4420 100644
--- a/internal/testsuite/testsuite.go
+++ b/internal/testsuite/testsuite.go
@@ -12,6 +12,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"iter"
 	"log"
 	"os"
 	"os/exec"
@@ -286,6 +287,20 @@ func Poll(
 	}
 }
 
+// PollPID polls until the process referred to by pid no longer exists.
+func PollPID(name string, pid int) {
+	for {
+		time.Sleep(10 * time.Millisecond)
+		if err := syscall.Kill(pid, 0); err == nil {
+			continue
+		} else if errors.Is(err, syscall.ESRCH) {
+			break
+		} else {
+			log.Fatalf("cannot poll for %s: %v", name, err)
+		}
+	}
+}
+
 const (
 	// XDGRuntimeDir is the hardcoded XDG runtime directory of the only user
 	// mapped in hsurc.
@@ -404,6 +419,12 @@ func MustStartSway(
 		[]string{SwayEnv},
 		"swaymsg",
 	)
+	Poll(
+		50*time.Millisecond,
+		cred,
+		[]string{"DISPLAY=:0"},
+		"xset", "q",
+	)
 	log.Printf("sway available via %s", SwayEnv)
 }
 
@@ -412,6 +433,82 @@ func TerminateSway(cred *syscall.Credential) {
 	MustFail(cred, []string{SwayEnv}, quietServices, "swaymsg", "exit")
 }
 
+// SwayNode represents a Sway window with hakurei-specific fields.
+type SwayNode struct {
+	// Window title.
+	Name string `json:"name"`
+	// Descendents.
+	Nodes []SwayNode `json:"nodes"`
+	// Wayland app id.
+	AppID string `json:"app_id,omitempty"`
+
+	// Should always be app.hakurei for cmd/hakurei containers.
+	Sandbox string `json:"sandbox_engine,omitempty"`
+	// Configured via [hst].
+	ID string `json:"sandbox_app_id,omitempty"`
+	// Hakurei instance identifier.
+	Instance hst.ID `json:"sandbox_instance_id,omitempty"`
+}
+
+// All returns an iterator over n and all descendents of n.
+func (n *SwayNode) All() iter.Seq[*SwayNode] {
+	return func(yield func(*SwayNode) bool) {
+		if !yield(n) {
+			return
+		}
+
+		for _, d := range n.Nodes {
+			for v := range d.All() {
+				if !yield(v) {
+					return
+				}
+			}
+		}
+	}
+}
+
+// PollSwayWindow continuously polls sway for a Wayland client to appear.
+func PollSwayWindow(
+	ctx context.Context,
+	cred *syscall.Credential,
+	f func(n *SwayNode) bool,
+) *SwayNode {
+	for {
+		time.Sleep(25 * time.Millisecond)
+		cmd := exec.CommandContext(ctx, "swaymsg", "-t", "get_tree")
+		cmd.Env = []string{SwayEnv}
+		cmd.SysProcAttr = &syscall.SysProcAttr{
+			Pdeathsig:  syscall.SIGKILL,
+			Credential: cred,
+		}
+		cmd.Stderr = os.Stderr
+
+		r, err := cmd.StdoutPipe()
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		if err = cmd.Start(); err != nil {
+			log.Fatalf("poll swaymsg: %v", err)
+		}
+
+		var t SwayNode
+		if err = json.NewDecoder(r).Decode(&t); err != nil {
+			log.Fatalf("decode swaymsg: %v", err)
+		}
+
+		if err = cmd.Wait(); err != nil {
+			log.Fatalf("wait for swaymsg: %v", err)
+		}
+
+		for n := range t.All() {
+			if f(n) {
+				return n
+			}
+		}
+	}
+}
+
 // MustStartPipeWire starts a PipeWire server that is never explicitly
 // terminated.  The test suite is terminated if the PipeWire server terminates.
 func MustStartPipeWire(cred *syscall.Credential, dbusEnv string) (version string) {