From 9c73682be8a8741fc9394f7bd26b794a2209b1cc Mon Sep 17 00:00:00 2001
From: Terry Howe <terrylhowe@gmail.com>
Date: Tue, 3 Dec 2024 18:47:56 -0700
Subject: [PATCH] fix: broken tty tests

Signed-off-by: Terry Howe <terrylhowe@gmail.com>
---
 .../internal/display/status/tty_console_test.go   | 14 ++++++++++----
 cmd/oras/root/cp.go                               |  7 ++++++-
 cmd/oras/root/cp_test.go                          | 15 +++------------
 3 files changed, 19 insertions(+), 17 deletions(-)

diff --git a/cmd/oras/internal/display/status/tty_console_test.go b/cmd/oras/internal/display/status/tty_console_test.go
index 9deb9a234..2a0efb7f9 100644
--- a/cmd/oras/internal/display/status/tty_console_test.go
+++ b/cmd/oras/internal/display/status/tty_console_test.go
@@ -127,8 +127,10 @@ func TestTTYCopyHandler_OnCopySkipped(t *testing.T) {
 		t.Errorf("OnCopySkipped() should not return an error: %v", err)
 	}
 
-	ch.StopTracking()
-	if err = testutils.MatchPty(pty, slave, "Exists", mockFetcher.OciImage.MediaType, strconv.FormatInt(mockFetcher.OciImage.Size, 10), "100%"); err != nil {
+	if err = ch.StopTracking(); err != nil {
+		t.Errorf("StopTracking() should not return an error: %v", err)
+	}
+	if err = testutils.MatchPty(pty, slave, "Exists", "oci-image", strconv.FormatInt(mockFetcher.OciImage.Size, 10), "100.00%"); err != nil {
 		t.Fatal(err)
 	}
 }
@@ -144,12 +146,14 @@ func TestTTYCopyHandler_PostCopy(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer ch.StopTracking()
 
 	if ch.PostCopy(ctx, bogus) == nil {
 		t.Error("PostCopy() should return an error")
 	}
 
+	if err = ch.StopTracking(); err != nil {
+		t.Errorf("StopTracking() should not return an error: %v", err)
+	}
 	if err = testutils.MatchPty(pty, slave, "\x1b[?25l\x1b7\x1b[0m"); err != nil {
 		t.Fatal(err)
 	}
@@ -166,12 +170,14 @@ func TestTTYCopyHandler_PreCopy(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer ch.StopTracking()
 
 	if err = ch.PreCopy(ctx, mockFetcher.OciImage); err != nil {
 		t.Errorf("PreCopy() should not return an error: %v", err)
 	}
 
+	if err = ch.StopTracking(); err != nil {
+		t.Errorf("StopTracking() should not return an error: %v", err)
+	}
 	if err = testutils.MatchPty(pty, slave, "\x1b[?25l\x1b7\x1b[0m"); err != nil {
 		t.Fatal(err)
 	}
diff --git a/cmd/oras/root/cp.go b/cmd/oras/root/cp.go
index 79a194b91..6e9a06d6c 100644
--- a/cmd/oras/root/cp.go
+++ b/cmd/oras/root/cp.go
@@ -174,7 +174,12 @@ func doCopy(ctx context.Context, copyHandler status.CopyHandler, src oras.ReadOn
 	if err != nil {
 		return desc, err
 	}
-	defer copyHandler.StopTracking()
+	defer func() {
+		stopErr := copyHandler.StopTracking()
+		if err == nil {
+			err = stopErr
+		}
+	}()
 	extendedCopyOptions.OnCopySkipped = copyHandler.OnCopySkipped
 	extendedCopyOptions.PreCopy = copyHandler.PreCopy
 	extendedCopyOptions.PostCopy = copyHandler.PostCopy
diff --git a/cmd/oras/root/cp_test.go b/cmd/oras/root/cp_test.go
index c6c88b6d6..1e9b19ea1 100644
--- a/cmd/oras/root/cp_test.go
+++ b/cmd/oras/root/cp_test.go
@@ -25,7 +25,6 @@ import (
 	"net/http/httptest"
 	"net/url"
 	"os"
-	"strings"
 	"testing"
 
 	"github.com/opencontainers/go-digest"
@@ -33,7 +32,6 @@ import (
 	"oras.land/oras-go/v2/content/memory"
 	"oras.land/oras-go/v2/registry/remote"
 	"oras.land/oras/cmd/oras/internal/display/status"
-	"oras.land/oras/cmd/oras/internal/output"
 	"oras.land/oras/internal/testutils"
 )
 
@@ -131,9 +129,7 @@ func Test_doCopy(t *testing.T) {
 	opts.TTY = slave
 	opts.From.Reference = memDesc.Digest.String()
 	dst := memory.New()
-	builder := &strings.Builder{}
-	printer := output.NewPrinter(builder, os.Stderr)
-	handler := status.NewTextCopyHandler(printer, dst)
+	handler := status.NewTTYCopyHandler(opts.TTY)
 	// test
 	_, err = doCopy(context.Background(), handler, memStore, dst, &opts)
 	if err != nil {
@@ -155,9 +151,7 @@ func Test_doCopy_skipped(t *testing.T) {
 	var opts copyOptions
 	opts.TTY = slave
 	opts.From.Reference = memDesc.Digest.String()
-	builder := &strings.Builder{}
-	printer := output.NewPrinter(builder, os.Stderr)
-	handler := status.NewTextCopyHandler(printer, memStore)
+	handler := status.NewTTYCopyHandler(opts.TTY)
 
 	// test
 	_, err = doCopy(context.Background(), handler, memStore, memStore, &opts)
@@ -191,10 +185,7 @@ func Test_doCopy_mounted(t *testing.T) {
 		t.Fatal(err)
 	}
 	to.PlainHTTP = true
-	builder := &strings.Builder{}
-	printer := output.NewPrinter(builder, os.Stderr)
-	printer.Verbose = true
-	handler := status.NewTextCopyHandler(printer, to)
+	handler := status.NewTTYCopyHandler(opts.TTY)
 
 	// test
 	_, err = doCopy(context.Background(), handler, from, to, &opts)