From 2e312b59aa7d43fe448cae52a2c5ddc63f272691 Mon Sep 17 00:00:00 2001
From: Ryan Nowak <nowakra@gmail.com>
Date: Sun, 5 Jan 2025 10:58:24 -0800
Subject: [PATCH] Implement fix for failures caused in test logging

Logging in tests has an intentional limitation. Any logging performed after the end of a test will cause the test to fail.

Unfortunately for this, this is a really hard limitation. We do a lot of asynchronous processing, and we can't always guarantee that processing is complete when a test completes. See https://github.com/radius-project/radius/actions/runs/12591118184/job/35093755888 for an example failure. This is especially bad because it's never a consistent repro.

This approach is based on the solution that Kubernetes uses: https://github.com/kubernetes/klog/pull/337/files#diff-9c4e6a71ada6fa5d0ee80e05296e5b418351025a21b2f9dce9595da245272313 . They wrote their own logging package, so they can make the fix inside of it. We're using logr (which was created by the same author), so we have to make the fix outside of the the logging package, but it's the same fundamental idea.

Signed-off-by: Ryan Nowak <nowakra@gmail.com>
---
 test/testcontext/testcontext.go | 44 ++++++++++++++++++++++++++++++++-
 1 file changed, 43 insertions(+), 1 deletion(-)

diff --git a/test/testcontext/testcontext.go b/test/testcontext/testcontext.go
index 6f6a65a330..999c41e9c3 100644
--- a/test/testcontext/testcontext.go
+++ b/test/testcontext/testcontext.go
@@ -18,6 +18,7 @@ package testcontext
 
 import (
 	"context"
+	"sync"
 	"testing"
 	"time"
 
@@ -48,8 +49,12 @@ func Wrap(t *testing.T, ctx context.Context) (context.Context, context.CancelFun
 		ctx = context.Background()
 	}
 
+	// See comments on testSink for why we need this.
+	wrapper := &testSink{wrapped: t}
+	t.Cleanup(wrapper.Stop)
+
 	// Setting verbosity so that everything gets logged.
-	ctx = logr.NewContext(ctx, testr.NewWithOptions(t, testr.Options{LogTimestamp: true, Verbosity: 10000}))
+	ctx = logr.NewContext(ctx, testr.NewWithInterface(wrapper, testr.Options{LogTimestamp: true, Verbosity: 10000}))
 	deadline, ok := t.Deadline()
 	if ok {
 		return context.WithDeadline(ctx, deadline)
@@ -57,3 +62,40 @@ func Wrap(t *testing.T, ctx context.Context) (context.Context, context.CancelFun
 		return context.WithCancel(ctx)
 	}
 }
+
+// testSink wraps testing.T for logging purposes.
+//
+// This type exists to work around an intentional limitation of testing.T.Log. It's not allowed to log after
+// the test completes.
+//
+// Unfortunately that's not always viable for us given the amount of async processing we do in Radius.
+//
+// This type can hook onto testing.T.Cleanup and no-op all logging when the test completes.
+type testSink struct {
+	wrapped *testing.T
+	mutex   sync.Mutex
+	stopped bool
+}
+
+var _ testr.TestingT = (*testSink)(nil)
+
+// Stop is called when the test terminates, this will disable logging.
+func (t *testSink) Stop() {
+	t.mutex.Lock()
+	defer t.mutex.Unlock()
+	t.stopped = true
+}
+
+// Helper implements testr.TestingT.
+func (t *testSink) Helper() {
+	if !t.stopped {
+		t.wrapped.Helper()
+	}
+}
+
+// Log implements testr.TestingT.
+func (t *testSink) Log(args ...any) {
+	if !t.stopped {
+		t.wrapped.Log(args...)
+	}
+}