Refactor tests to handle batch message returns

main_test.go

@@ -2,8 +2,6 @@ package main
 
 import (
 	"fmt"
-	"reflect"
-	"runtime"
 	"strings"
 	"testing"
 	"time"
@@ -11,6 +9,12 @@ import (
 	tea "github.com/charmbracelet/bubbletea"
 )
 
+func notThis[T comparable](f compareFunc[T]) compareFunc[T] {
+	return func(a, b T) bool {
+		return !f(a, b)
+	}
+}
+
 type compareFunc[T comparable] func(T, T) bool
 
 func assertFunc[T comparable](t *testing.T, compare compareFunc[T], actual, expected T, msg string) {
@@ -71,8 +75,8 @@ func TestRunCommands(t *testing.T) {
 	m.startCommands(m.onFocusStart)
 }
 
-func keyMsgs(keys ...interface{}) []tea.Msg {
+func keyMsgs(keys ...interface{}) []tea.KeyMsg {
-	keyMessages := []tea.Msg{}
+	keyMessages := []tea.KeyMsg{}
 
 	for _, key := range keys {
 		switch keyType := key.(type) {
@@ -90,105 +94,214 @@ func keyMsgs(keys ...interface{}) []tea.Msg {
 	return keyMessages
 }
 
-func sendKeys(m model, keys ...interface{}) (model, tea.Cmd) {
+func sendKeys[T tea.Model](m T, keys ...interface{}) (T, []tea.Cmd) {
 	var updatedModel tea.Model
 
 	var cmd tea.Cmd
 
+	var cmds []tea.Cmd
+
 	for _, key := range keyMsgs(keys...) {
 		updatedModel, cmd = m.Update(key)
-		m = updatedModel.(model)
+		m = updatedModel.(T)
 
 		if cmd != nil {
-			updatedModel, cmd = m.Update(cmd())
+			m, cmds, _ = batchCmdUpdate(m, cmd)
-			m = updatedModel.(model)
 		}
 	}
 
-	return m, cmd
+	return m, cmds
+}
+
+func batchCmdUpdate[T tea.Model](m T, cmd tea.Cmd) (T, []tea.Cmd, error) {
+	var updatedModel tea.Model
+
+	var resultCmd tea.Cmd
+
+	var cmds []tea.Cmd
+
+	resultMsg := cmd()
+
+	batchMsg, ok := resultMsg.(tea.BatchMsg)
+	if !ok {
+		updatedModel, resultCmd = m.Update(resultMsg)
+		m = updatedModel.(T)
+
+		cmds = append(cmds, resultCmd)
+
+		return m, cmds, fmt.Errorf("Expected a batch message, got %v", resultMsg)
+	}
+
+	// Apply batch messages
+	for _, cmd := range batchMsg {
+		updatedModel, resultCmd = m.Update(cmd())
+		m = updatedModel.(T)
+
+		cmds = append(cmds, resultCmd)
+	}
+
+	return m, cmds, nil
+}
+
+type mockModel struct {
+	messages []tea.Msg
+}
+
+func (m mockModel) Init() tea.Cmd {
+	return nil
+}
+
+func (m mockModel) View() string {
+	s := ""
+
+	for _, msg := range m.messages {
+		if keyMsg, ok := msg.(tea.KeyMsg); ok {
+			if keyMsg.Type == tea.KeyRunes {
+				s += string(keyMsg.Runes)
+			}
+		}
+	}
+
+	return s
+}
+
+func (m mockModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
+	var cmds []tea.Cmd
+
+	m.messages = append(m.messages, msg)
+
+	if msg, ok := msg.(tea.KeyMsg); ok {
+		if msg.Type == tea.KeyEnter {
+			cmds = append(cmds, func() tea.Msg { return tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'\n'}} })
+		}
+	}
+
+	return m, tea.Batch(cmds...)
+}
+
+func TestBatchCmdUpdate(t *testing.T) {
+	m := mockModel{}
+
+	cmds := []tea.Cmd{
+		func() tea.Msg { return tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'a'}} },
+		func() tea.Msg { return tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'b'}} },
+	}
+
+	m, _, err := batchCmdUpdate(m, tea.Batch(cmds...))
+	if err != nil {
+		t.Fatalf("Expected no error, got %v", err)
+	}
+
+	if len(m.messages) != 2 {
+		t.Fatalf("Expected 2 messages, got %v", len(m.messages))
+	}
+
+	if m.View() != "ab" {
+		t.Fatalf("Expected 'ab', got %s, %+v", m.View(), m)
+	}
+
+	m, _, err = batchCmdUpdate(m, func() tea.Msg { return tea.KeyMsg{Type: tea.KeyRunes, Runes: []rune{'c'}} })
+	if err == nil {
+		t.Fatalf("Expected error, got nil")
+	}
+
+	if len(m.messages) != 3 {
+		t.Fatalf("Expected 3 messages, got %v", len(m.messages))
+	}
+
+	if m.View() != "abc" {
+		t.Fatalf("Expected 'abc', got %v", m.View())
+	}
+}
+
+func TestSendKeys(t *testing.T) {
+	m := mockModel{}
+
+	m, _ = sendKeys(m, "ab")
+
+	if len(m.messages) != 2 {
+		t.Fatalf("Expected 2 messages, got %v", len(m.messages))
+	}
+
+	if m.View() != "ab" {
+		t.Fatalf("Expected 'ab', got %s", m.View())
+	}
+
+	m, _ = sendKeys(m, "c", tea.KeyEnter)
+
+	// We expect 5 because enter key should create a new message for the newline
+	if len(m.messages) != 5 {
+		t.Fatalf("Expected 5 messages, got %v", len(m.messages))
+	}
+
+	if m.View() != "abc\n" {
+		t.Fatalf("Expected 'abc\n', got %s", m.View())
+	}
 }
 
 // TestInputView tests the Update method of the model for the input view
 func TestInputView(t *testing.T) {
 	m := newModel(false, "#ffdd57", "#57ddff", 0, 0, 0, []string{}, []string{}, []string{})
 
-	var updatedModel tea.Model
+	var err error
 
-	var resultCmd tea.Cmd
+	m, _, err = batchCmdUpdate(m, m.Init())
-
+	if err != nil {
-	updatedModel, _ = m.Update(m.Init())
+		t.Fatalf("Expected batch command after init: %v", err)
-	m = updatedModel.(model)
+	}
 
 	// Verify we're on the input screen
-	assertEqual(t, m.currentScreen, INPUT_SCREEN, "Expected currentScreen to be inputScreen")
+	assertEqual(t, m.currentScreen, INPUT_SCREEN, "Expected currentScreen to be INPUT_SCREEN")
-	assertFunc(t, strings.Contains, m.View(), "Break time", "Expected view to contain 'enter next'")
+	assertFunc(t, strings.Contains, m.View(), "Break time", "Expected view to contain 'Break time'")
 
 	// Fill the form
-	m, resultCmd = sendKeys(m, "10m", tea.KeyEnter, "10m", tea.KeyEnter, "2", tea.KeyEnter)
+	m, _ = sendKeys(m, "10m", tea.KeyEnter, "10m", tea.KeyEnter, "2", tea.KeyEnter)
-	assertNotEqual(t, resultCmd, nil, "Expected resultCmd to be not nil")
 
-	// Apply final next result
+	// Pass on submit command, doesn't happen because we have multiple layers of batching at the end
-	updatedModel, resultCmd = m.Update(resultCmd())
+	m, _, _ = batchCmdUpdate(m, formSubmit)
-	m = updatedModel.(model)
 
-	// Make sure the next command is to submit the form
+	assertEqual(t, m.focusTime, 10*time.Minute, "Expected focus time to be 10 minutes")
-	assertNotEqual(t, resultCmd, nil, "Expected resultCmd to be not nil")
+	assertEqual(t, m.breakTime, 10*time.Minute, "Expected break time to be 10 minutes")
-	assertEqual(
+	assertEqual(t, m.intervals, 2, "Expected rounds to be 2")
-		t,
-		reflect.ValueOf(resultCmd).Pointer(),
-		reflect.ValueOf(formSubmit).Pointer(),
-		fmt.Sprintf("Expected resultCmd to be formSubmit, found %v", runtime.FuncForPC(reflect.ValueOf(resultCmd).Pointer()).Name()),
-	)
 
-	// Apply submit form command
+	// Pass on start command, doesn't happen because we have multiple layers of batching at the end
-	updatedModel, resultCmd = m.Update(resultCmd())
+	m, _, _ = batchCmdUpdate(m, startTimer)
-	m = updatedModel.(model)
 
-	assertEqual(t, m.err, nil, "Expected no error")
+	assertFunc(t, notThis(strings.Contains), m.View(), "Break time", "Expected view to NOT contain 'Break time'")
-	assertNotEqual(t, resultCmd, nil, "Expected resultCmd to be not nil")
+	assertEqual(t, m.currentScreen, TIMER_SCREEN, "Expected currentScreen to be TIMER_SCREEN")
-	assertEqual(
-		t,
-		reflect.ValueOf(resultCmd).Pointer(),
-		reflect.ValueOf(startTimer).Pointer(),
-		fmt.Sprintf("Expected resultCmd to be startTimer, found %v", runtime.FuncForPC(reflect.ValueOf(resultCmd).Pointer()).Name()),
-	)
 }
 
 func TestTimerView(t *testing.T) {
 	m := newModel(false, "#ffdd57", "#57ddff", 10*time.Minute, 10*time.Minute, 2, []string{}, []string{}, []string{})
 
-	var updatedModel tea.Model
-
 	// Init model with timer values
-	updatedModel, _ = m.Update(m.Init())
+	m, _, err := batchCmdUpdate(m, m.Init())
-	m = updatedModel.(model)
+	if err != nil {
+		t.Fatalf("Expected batch command after init: %v", err)
+	}
 
 	// Start timer (batch result from above doesn't apply here)
-	updatedModel, _ = m.Update(startTimer())
+	m, _, _ = batchCmdUpdate(m, startTimer)
-	m = updatedModel.(model)
 
 	// Test timer view
 	assertEqual(t, m.currentScreen, TIMER_SCREEN, "Expected currentScreen to be timerScreen")
 	assertFunc(t, strings.Contains, m.View(), "Focus", "Expected view to contain 'Focus'")
 	assertEqual(t, m.state, "Focus", "Expected state to be 'Focus'")
 
-	oneSec := timeMsg(time.Now().Add(1 * time.Second))
+	oneSecCmd := func() tea.Msg { return timeMsg(time.Now().Add(1 * time.Second)) }
-	updatedModel, _ = m.Update(oneSec)
+	m, _, _ = batchCmdUpdate(m, oneSecCmd)
-	m = updatedModel.(model)
 
 	assertEqual(t, m.state, "Focus", "Expected state to be 'Focus'")
 
 	// Test switch to break time
 	m.startTime = m.startTime.Add(-10 * time.Minute)
-	updatedModel, _ = m.Update(oneSec)
+	m, _, _ = batchCmdUpdate(m, oneSecCmd)
-	m = updatedModel.(model)
 
 	assertEqual(t, m.state, "Break", "Expected state to be 'Break'")
 
 	// Switch back to focus time
 	m.startTime = m.startTime.Add(-10 * time.Minute)
-	updatedModel, _ = m.Update(oneSec)
+	m, _, _ = batchCmdUpdate(m, oneSecCmd)
-	m = updatedModel.(model)
 
 	assertEqual(t, m.state, "Focus", "Expected state to be 'Focus'")
 	t.Logf("Incorrect state %+v", m)