From d6a19f0eb5a983610bd65a1647f5955abe3ee69e Mon Sep 17 00:00:00 2001
From: Ayman Bagabas <ayman.bagabas@gmail.com>
Date: Mon, 12 Aug 2024 16:59:45 -0400
Subject: [PATCH] fix: wrap ErrProgramKilled error

This gives the caller more context on _why_ the program was killed by
wrapping the error with the context error.
---
 tea.go      |  2 +-
 tea_test.go | 10 +++++++---
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/tea.go b/tea.go
index 509ee91552..62cd6415bb 100644
--- a/tea.go
+++ b/tea.go
@@ -582,7 +582,7 @@ func (p *Program) Run() (Model, error) {
 	model, err := p.eventLoop(model, cmds)
 	killed := p.ctx.Err() != nil
 	if killed {
-		err = ErrProgramKilled
+		err = fmt.Errorf("%w: %s", ErrProgramKilled, p.ctx.Err())
 	} else {
 		// Ensure we rendered the final state of the model.
 		p.renderer.write(model.View())
diff --git a/tea_test.go b/tea_test.go
index d65f0d7fe3..981851be0b 100644
--- a/tea_test.go
+++ b/tea_test.go
@@ -3,6 +3,7 @@ package tea
 import (
 	"bytes"
 	"context"
+	"errors"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -46,7 +47,10 @@ func TestTeaModel(t *testing.T) {
 	var in bytes.Buffer
 	in.Write([]byte("q"))
 
-	p := NewProgram(&testModel{}, WithInput(&in), WithOutput(&buf))
+	ctx, cancel := context.WithTimeout(context.TODO(), 3*time.Second)
+	defer cancel()
+
+	p := NewProgram(&testModel{}, WithInput(&in), WithOutput(&buf), WithContext(ctx))
 	if _, err := p.Run(); err != nil {
 		t.Fatal(err)
 	}
@@ -134,7 +138,7 @@ func TestTeaKill(t *testing.T) {
 		}
 	}()
 
-	if _, err := p.Run(); err != ErrProgramKilled {
+	if _, err := p.Run(); !errors.Is(err, ErrProgramKilled) {
 		t.Fatalf("Expected %v, got %v", ErrProgramKilled, err)
 	}
 }
@@ -156,7 +160,7 @@ func TestTeaContext(t *testing.T) {
 		}
 	}()
 
-	if _, err := p.Run(); err != ErrProgramKilled {
+	if _, err := p.Run(); !errors.Is(err, ErrProgramKilled) {
 		t.Fatalf("Expected %v, got %v", ErrProgramKilled, err)
 	}
 }