From 5d1ffa74cdae270160b8558ac1012aa35f267df6 Mon Sep 17 00:00:00 2001
From: Christian Muehlhaeuser <muesli@gmail.com>
Date: Fri, 23 Sep 2022 00:40:18 +0200
Subject: [PATCH] test: quit/kill model after the first render

---
 tea_test.go | 55 +++++++++++++++++++++++++++++++++++++++++++++++++----
 tty.go      |  2 +-
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/tea_test.go b/tea_test.go
index 3c35074..7a3d315 100644
--- a/tea_test.go
+++ b/tea_test.go
@@ -2,16 +2,20 @@ package tea
 
 import (
 	"bytes"
+	"sync/atomic"
 	"testing"
+	"time"
 )
 
-type testModel struct{}
+type testModel struct {
+	executed atomic.Value
+}
 
 func (m testModel) Init() Cmd {
 	return nil
 }
 
-func (m testModel) Update(msg Msg) (Model, Cmd) {
+func (m *testModel) Update(msg Msg) (Model, Cmd) {
 	switch msg.(type) {
 	case KeyMsg:
 		return m, Quit
@@ -19,7 +23,8 @@ func (m testModel) Update(msg Msg) (Model, Cmd) {
 	return m, nil
 }
 
-func (m testModel) View() string {
+func (m *testModel) View() string {
+	m.executed.Store(true)
 	return "success\n"
 }
 
@@ -28,7 +33,7 @@ func TestTeaModel(t *testing.T) {
 	var in bytes.Buffer
 	in.Write([]byte("q"))
 
-	p := NewProgram(testModel{}, WithInput(&in), WithOutput(&buf))
+	p := NewProgram(&testModel{}, WithInput(&in), WithOutput(&buf))
 	if err := p.Start(); err != nil {
 		t.Fatal(err)
 	}
@@ -37,3 +42,45 @@ func TestTeaModel(t *testing.T) {
 		t.Fatal("no output")
 	}
 }
+
+func TestTeaQuit(t *testing.T) {
+	var buf bytes.Buffer
+	var in bytes.Buffer
+
+	m := &testModel{}
+	p := NewProgram(m, WithInput(&in), WithOutput(&buf))
+	go func() {
+		for {
+			time.Sleep(time.Millisecond)
+			if m.executed.Load() != nil {
+				p.Quit()
+				return
+			}
+		}
+	}()
+
+	if err := p.Start(); err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestTeaKill(t *testing.T) {
+	var buf bytes.Buffer
+	var in bytes.Buffer
+
+	m := &testModel{}
+	p := NewProgram(m, WithInput(&in), WithOutput(&buf))
+	go func() {
+		for {
+			time.Sleep(time.Millisecond)
+			if m.executed.Load() != nil {
+				p.Kill()
+				return
+			}
+		}
+	}()
+
+	if err := p.Start(); err != nil {
+		t.Fatal(err)
+	}
+}
diff --git a/tty.go b/tty.go
index e4f75cf..5259638 100644
--- a/tty.go
+++ b/tty.go
@@ -26,7 +26,7 @@ func (p *Program) initTerminal() error {
 
 // restoreTerminalState restores the terminal to the state prior to running the
 // Bubble Tea program.
-func (p Program) restoreTerminalState() error {
+func (p *Program) restoreTerminalState() error {
 	p.output.ShowCursor()
 
 	if p.console != nil {