Extract generic drain test helper for screen adapters

Amolith created

Change summary

internal/ui/screens/drain_test.go      | 31 ++++++++++++++++++++
internal/ui/screens/filepicker_test.go | 28 +----------------
internal/ui/screens/overwrite_test.go  | 32 +++-----------------
internal/ui/screens/preset_test.go     | 43 +++++----------------------
internal/ui/screens/snapshot_test.go   | 42 ++++++--------------------
internal/ui/screens/target_test.go     | 37 +++--------------------
6 files changed, 63 insertions(+), 150 deletions(-)

Detailed changes

internal/ui/screens/drain_test.go 🔗

@@ -0,0 +1,31 @@
+package screens
+
+import (
+	tea "charm.land/bubbletea/v2"
+
+	"git.secluded.site/keld/internal/ui"
+)
+
+// drain feeds commands back into a screen until a DoneMsg or BackMsg
+// is produced, or the command chain is exhausted. This is necessary
+// because huh uses internal message chains (nextFieldMsg →
+// nextGroupMsg → StateCompleted) that must be processed sequentially.
+func drain[S ui.Screen](s S, initialCmd tea.Cmd) (S, tea.Cmd) {
+	cmd := initialCmd
+	for cmd != nil {
+		msg := cmd()
+		if msg == nil {
+			return s, nil
+		}
+		switch msg.(type) {
+		case ui.DoneMsg:
+			return s, cmd
+		case ui.BackMsg:
+			return s, cmd
+		}
+		var screen ui.Screen
+		screen, cmd = s.Update(msg)
+		s = screen.(S)
+	}
+	return s, nil
+}

internal/ui/screens/filepicker_test.go 🔗

@@ -23,35 +23,13 @@ func testNodes() []restic.LsNode {
 	}
 }
 
-// drainFilePicker feeds commands back into the file picker screen
-// until a DoneMsg or BackMsg is produced, or the chain is exhausted.
-func drainFilePicker(fp *FilePicker, initialCmd tea.Cmd) (*FilePicker, tea.Cmd) {
-	cmd := initialCmd
-	for cmd != nil {
-		msg := cmd()
-		if msg == nil {
-			return fp, nil
-		}
-		switch msg.(type) {
-		case ui.DoneMsg:
-			return fp, cmd
-		case ui.BackMsg:
-			return fp, cmd
-		}
-		var screen ui.Screen
-		screen, cmd = fp.Update(msg)
-		fp = screen.(*FilePicker)
-	}
-	return fp, nil
-}
-
 // simulateFileLoad sends a filesLoadedMsg directly to the screen.
 func simulateFileLoad(fp *FilePicker, nodes []restic.LsNode, err error) *FilePicker {
 	gen := atomic.LoadInt64(&fp.gen)
 	msg := filesLoadedMsg{gen: gen, nodes: nodes, err: err}
 	screen, cmd := fp.Update(msg)
 	fp = screen.(*FilePicker)
-	fp, _ = drainFilePicker(fp, cmd)
+	fp, _ = drain(fp, cmd)
 	return fp
 }
 
@@ -304,7 +282,7 @@ func TestFilePickerIncludesWithPartialSelection(t *testing.T) {
 	// (should be dir1, sorted dirs first) via space.
 	screen, cmd := fp.Update(tea.KeyPressMsg{Code: ' ', Text: " "})
 	fp = screen.(*FilePicker)
-	fp, _ = drainFilePicker(fp, cmd)
+	fp, _ = drain(fp, cmd)
 
 	// Confirm with enter.
 	screen, cmd = fp.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
@@ -341,7 +319,7 @@ func TestFilePickerSelectionShowsCount(t *testing.T) {
 	// Toggle first item and confirm.
 	screen, cmd := fp.Update(tea.KeyPressMsg{Code: ' ', Text: " "})
 	fp = screen.(*FilePicker)
-	fp, _ = drainFilePicker(fp, cmd)
+	fp, _ = drain(fp, cmd)
 
 	screen, _ = fp.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
 	fp = screen.(*FilePicker)

internal/ui/screens/overwrite_test.go 🔗

@@ -8,28 +8,6 @@ import (
 	"git.secluded.site/keld/internal/ui"
 )
 
-// drainOverwrite feeds commands back into the overwrite screen until
-// a DoneMsg or BackMsg is produced, or the command chain is exhausted.
-func drainOverwrite(o *Overwrite, initialCmd tea.Cmd) (*Overwrite, tea.Cmd) {
-	cmd := initialCmd
-	for cmd != nil {
-		msg := cmd()
-		if msg == nil {
-			return o, nil
-		}
-		switch msg.(type) {
-		case ui.DoneMsg:
-			return o, cmd
-		case ui.BackMsg:
-			return o, cmd
-		}
-		var screen ui.Screen
-		screen, cmd = o.Update(msg)
-		o = screen.(*Overwrite)
-	}
-	return o, nil
-}
-
 func TestOverwriteTitle(t *testing.T) {
 	t.Parallel()
 
@@ -88,7 +66,7 @@ func TestOverwriteCompleteReturnsDone(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = o.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	_, cmd = drainOverwrite(screen.(*Overwrite), cmd)
+	_, cmd = drain(screen.(*Overwrite), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after form completion")
@@ -110,7 +88,7 @@ func TestOverwriteDefaultIsIfChanged(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = o.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	o, _ = drainOverwrite(screen.(*Overwrite), cmd)
+	o, _ = drain(screen.(*Overwrite), cmd)
 
 	if got := o.Value(); got != "if-changed" {
 		t.Errorf("Value() = %q, want %q (default should be if-changed)", got, "if-changed")
@@ -128,7 +106,7 @@ func TestOverwriteSelectionShowsValue(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = o.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	o, _ = drainOverwrite(screen.(*Overwrite), cmd)
+	o, _ = drain(screen.(*Overwrite), cmd)
 
 	if got := o.Selection(); got != "if-changed" {
 		t.Errorf("Selection() = %q, want %q", got, "if-changed")
@@ -146,7 +124,7 @@ func TestOverwriteReinitAfterBack(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = o.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	o, _ = drainOverwrite(screen.(*Overwrite), cmd)
+	o, _ = drain(screen.(*Overwrite), cmd)
 
 	// Re-init (simulates back navigation).
 	o.Init()
@@ -154,7 +132,7 @@ func TestOverwriteReinitAfterBack(t *testing.T) {
 
 	// Should be functional again.
 	screen, cmd = o.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	_, cmd = drainOverwrite(screen.(*Overwrite), cmd)
+	_, cmd = drain(screen.(*Overwrite), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after re-init")

internal/ui/screens/preset_test.go 🔗

@@ -14,33 +14,6 @@ func testPresets() []string {
 	return []string{"home@cloud", "work@local", "media"}
 }
 
-// drainPreset feeds commands back into the preset screen until a
-// DoneMsg or BackMsg is produced, or the command chain is exhausted.
-// This is necessary because huh uses internal message chains
-// (nextFieldMsg → nextGroupMsg → StateCompleted) that must be
-// processed sequentially.
-func drainPreset(p *Preset, initialCmd tea.Cmd) (*Preset, tea.Cmd) {
-	cmd := initialCmd
-	for cmd != nil {
-		msg := cmd()
-		if msg == nil {
-			return p, nil
-		}
-		// Check if this is one of our terminal messages.
-		switch msg.(type) {
-		case ui.DoneMsg:
-			return p, cmd
-		case ui.BackMsg:
-			return p, cmd
-		}
-		// Otherwise feed it back to the screen.
-		var screen ui.Screen
-		screen, cmd = p.Update(msg)
-		p = screen.(*Preset)
-	}
-	return p, nil
-}
-
 func TestPresetTitle(t *testing.T) {
 	t.Parallel()
 
@@ -132,7 +105,7 @@ func TestPresetCompleteReturnsDone(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	_, cmd = drainPreset(screen.(*Preset), cmd)
+	_, cmd = drain(screen.(*Preset), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after form completion")
@@ -155,7 +128,7 @@ func TestPresetSelectionValue(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	p, _ = drainPreset(screen.(*Preset), cmd)
+	p, _ = drain(screen.(*Preset), cmd)
 
 	if got := p.Selection(); got != "home@cloud" {
 		t.Errorf("Selection() = %q, want %q", got, "home@cloud")
@@ -175,7 +148,7 @@ func TestPresetGlobalDefaultDisplayLabel(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	p, _ = drainPreset(screen.(*Preset), cmd)
+	p, _ = drain(screen.(*Preset), cmd)
 
 	// Selection() should return a display-friendly label, not "".
 	sel := p.Selection()
@@ -196,7 +169,7 @@ func TestPresetReinitPreservesSelection(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	p, _ = drainPreset(screen.(*Preset), cmd)
+	p, _ = drain(screen.(*Preset), cmd)
 
 	if p.Selection() != "home@cloud" {
 		t.Fatalf("precondition: Selection() = %q, want %q", p.Selection(), "home@cloud")
@@ -223,7 +196,7 @@ func TestPresetReinitAfterBack(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	p, _ = drainPreset(screen.(*Preset), cmd)
+	p, _ = drain(screen.(*Preset), cmd)
 
 	// Re-initialise (simulates the session navigating back to this screen).
 	p.Init()
@@ -231,7 +204,7 @@ func TestPresetReinitAfterBack(t *testing.T) {
 
 	// Should be functional again — selecting should produce DoneCmd.
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	_, cmd = drainPreset(screen.(*Preset), cmd)
+	_, cmd = drain(screen.(*Preset), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after re-init, form may not have been reset")
@@ -276,7 +249,7 @@ func TestPresetValueReturnsActualPresetName(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	p, _ = drainPreset(screen.(*Preset), cmd)
+	p, _ = drain(screen.(*Preset), cmd)
 
 	if got := p.Value(); got != "home@cloud" {
 		t.Errorf("Value() = %q, want %q", got, "home@cloud")
@@ -294,7 +267,7 @@ func TestPresetValueEmptyForGlobalDefaults(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = p.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	p, _ = drainPreset(screen.(*Preset), cmd)
+	p, _ = drain(screen.(*Preset), cmd)
 
 	if got := p.Value(); got != "" {
 		t.Errorf("Value() = %q, want empty for global defaults", got)

internal/ui/screens/snapshot_test.go 🔗

@@ -32,28 +32,6 @@ func testSnapshots() []restic.Snapshot {
 	}
 }
 
-// drainSnapshot feeds commands back into the snapshot screen until a
-// DoneMsg or BackMsg is produced, or the command chain is exhausted.
-func drainSnapshot(s *Snapshot, initialCmd tea.Cmd) (*Snapshot, tea.Cmd) {
-	cmd := initialCmd
-	for cmd != nil {
-		msg := cmd()
-		if msg == nil {
-			return s, nil
-		}
-		switch msg.(type) {
-		case ui.DoneMsg:
-			return s, cmd
-		case ui.BackMsg:
-			return s, cmd
-		}
-		var screen ui.Screen
-		screen, cmd = s.Update(msg)
-		s = screen.(*Snapshot)
-	}
-	return s, nil
-}
-
 // simulateLoad sends a snapshotsLoadedMsg directly to the screen,
 // bypassing tea.Batch. This is how tests deliver async results —
 // in the real runtime, the Batch command produces this message.
@@ -63,7 +41,7 @@ func simulateLoad(s *Snapshot, snaps []restic.Snapshot, err error) *Snapshot {
 	screen, cmd := s.Update(msg)
 	s = screen.(*Snapshot)
 	// Drain any form init commands.
-	s, _ = drainSnapshot(s, cmd)
+	s, _ = drain(s, cmd)
 	return s
 }
 
@@ -120,7 +98,7 @@ func TestSnapshotLoadsAndPresentsSelect(t *testing.T) {
 
 	// Should now be in the selecting phase. Press enter to pick first.
 	screen, cmd := s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	s, cmd = drainSnapshot(screen.(*Snapshot), cmd)
+	s, cmd = drain(screen.(*Snapshot), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after selecting a snapshot")
@@ -146,7 +124,7 @@ func TestSnapshotSelectionShowsShortID(t *testing.T) {
 
 	// Select first snapshot.
 	screen, cmd := s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	s, _ = drainSnapshot(screen.(*Snapshot), cmd)
+	s, _ = drain(screen.(*Snapshot), cmd)
 
 	if got := s.Selection(); got != "abcdef01" {
 		t.Errorf("Selection() = %q, want %q", got, "abcdef01")
@@ -175,7 +153,7 @@ func TestSnapshotErrorFallsBackToManual(t *testing.T) {
 	}
 
 	screen, cmd := s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	s, cmd = drainSnapshot(screen.(*Snapshot), cmd)
+	s, cmd = drain(screen.(*Snapshot), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after manual entry")
@@ -209,7 +187,7 @@ func TestSnapshotNoRepoFallsBackSilently(t *testing.T) {
 		s = screen.(*Snapshot)
 	}
 	screen, cmd := s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	s, cmd = drainSnapshot(screen.(*Snapshot), cmd)
+	s, cmd = drain(screen.(*Snapshot), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after manual entry")
@@ -243,7 +221,7 @@ func TestSnapshotNilLoaderGoesDirectToManual(t *testing.T) {
 	s = screen.(*Snapshot)
 
 	// Drain any init commands from the manual form.
-	s, _ = drainSnapshot(s, cmd)
+	s, _ = drain(s, cmd)
 
 	// With no loader, should start in manual entry immediately.
 	for _, ch := range "latest" {
@@ -251,7 +229,7 @@ func TestSnapshotNilLoaderGoesDirectToManual(t *testing.T) {
 		s = screen.(*Snapshot)
 	}
 	screen, cmd = s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	s, cmd = drainSnapshot(screen.(*Snapshot), cmd)
+	s, cmd = drain(screen.(*Snapshot), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after manual entry")
@@ -280,7 +258,7 @@ func TestSnapshotManualEntryFromSelect(t *testing.T) {
 
 	// Select the manual entry option.
 	screen, cmd := s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	s, _ = drainSnapshot(screen.(*Snapshot), cmd)
+	s, _ = drain(screen.(*Snapshot), cmd)
 
 	// Should now be in manual entry phase. Type an ID.
 	for _, ch := range "abc123:subfolder" {
@@ -288,7 +266,7 @@ func TestSnapshotManualEntryFromSelect(t *testing.T) {
 		s = screen.(*Snapshot)
 	}
 	screen, cmd = s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	s, cmd = drainSnapshot(screen.(*Snapshot), cmd)
+	s, cmd = drain(screen.(*Snapshot), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after manual entry from select")
@@ -364,7 +342,7 @@ func TestSnapshotStaleLoadIgnored(t *testing.T) {
 
 	// Now should be in selecting phase.
 	screen, cmd := s.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	_, cmd = drainSnapshot(screen.(*Snapshot), cmd)
+	_, cmd = drain(screen.(*Snapshot), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after selecting from fresh load")

internal/ui/screens/target_test.go 🔗

@@ -8,31 +8,6 @@ import (
 	"git.secluded.site/keld/internal/ui"
 )
 
-// drainTarget feeds commands back into the target screen until a
-// DoneMsg or BackMsg is produced, or the command chain is exhausted.
-// This is necessary because huh uses internal message chains
-// (nextFieldMsg → nextGroupMsg → StateCompleted) that must be
-// processed sequentially.
-func drainTarget(t *Target, initialCmd tea.Cmd) (*Target, tea.Cmd) {
-	cmd := initialCmd
-	for cmd != nil {
-		msg := cmd()
-		if msg == nil {
-			return t, nil
-		}
-		switch msg.(type) {
-		case ui.DoneMsg:
-			return t, cmd
-		case ui.BackMsg:
-			return t, cmd
-		}
-		var screen ui.Screen
-		screen, cmd = t.Update(msg)
-		t = screen.(*Target)
-	}
-	return t, nil
-}
-
 func TestTargetTitle(t *testing.T) {
 	t.Parallel()
 
@@ -95,7 +70,7 @@ func TestTargetCompleteReturnsDone(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = tgt.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	_, cmd = drainTarget(screen.(*Target), cmd)
+	_, cmd = drain(screen.(*Target), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after form completion")
@@ -120,7 +95,7 @@ func TestTargetSelectionShowsPath(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = tgt.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	tgt, _ = drainTarget(screen.(*Target), cmd)
+	tgt, _ = drain(screen.(*Target), cmd)
 
 	if got := tgt.Selection(); got != "/tmp/restore" {
 		t.Errorf("Selection() = %q, want %q", got, "/tmp/restore")
@@ -141,7 +116,7 @@ func TestTargetValueReturnsPath(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = tgt.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	tgt, _ = drainTarget(screen.(*Target), cmd)
+	tgt, _ = drain(screen.(*Target), cmd)
 
 	if got := tgt.Value(); got != "/tmp/restore" {
 		t.Errorf("Value() = %q, want %q", got, "/tmp/restore")
@@ -162,7 +137,7 @@ func TestTargetReinitAfterBack(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = tgt.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	tgt, _ = drainTarget(screen.(*Target), cmd)
+	tgt, _ = drain(screen.(*Target), cmd)
 
 	// Re-init (simulates back navigation).
 	tgt.Init()
@@ -173,7 +148,7 @@ func TestTargetReinitAfterBack(t *testing.T) {
 		tgt.Update(tea.KeyPressMsg{Code: ch, Text: string(ch)})
 	}
 	screen, cmd = tgt.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	_, cmd = drainTarget(screen.(*Target), cmd)
+	_, cmd = drain(screen.(*Target), cmd)
 
 	if cmd == nil {
 		t.Fatal("expected DoneCmd after re-init")
@@ -197,7 +172,7 @@ func TestTargetPreservesValueOnReinit(t *testing.T) {
 	var screen ui.Screen
 	var cmd tea.Cmd
 	screen, cmd = tgt.Update(tea.KeyPressMsg{Code: tea.KeyEnter})
-	tgt, _ = drainTarget(screen.(*Target), cmd)
+	tgt, _ = drain(screen.(*Target), cmd)
 
 	// Re-init should preserve the entered value.
 	tgt.Init()