From cafbbcd8dee53918b25c8b25afd9e4f728d121b9 Mon Sep 17 00:00:00 2001
From: Jack McCluskey <thejackmccluskey@gmail.com>
Date: Tue, 12 Sep 2023 14:50:26 -0400
Subject: [PATCH] Address comments

---
 sdks/go/container/tools/buffered_logging.go   |  7 +++--
 .../container/tools/buffered_logging_test.go  | 28 +++++++++++++++----
 sdks/python/container/piputil.go              |  2 +-
 3 files changed, 27 insertions(+), 10 deletions(-)

diff --git a/sdks/go/container/tools/buffered_logging.go b/sdks/go/container/tools/buffered_logging.go
index e31a79eb154d..445d19fabfdc 100644
--- a/sdks/go/container/tools/buffered_logging.go
+++ b/sdks/go/container/tools/buffered_logging.go
@@ -36,17 +36,18 @@ type BufferedLogger struct {
 	lastFlush            time.Time
 	flushInterval        time.Duration
 	periodicFlushContext context.Context
+	now                  func() time.Time
 }
 
 // NewBufferedLogger returns a new BufferedLogger type by reference.
 func NewBufferedLogger(logger *Logger) *BufferedLogger {
-	return &BufferedLogger{logger: logger, lastFlush: time.Now(), flushInterval: time.Duration(math.MaxInt64), periodicFlushContext: context.Background()}
+	return &BufferedLogger{logger: logger, lastFlush: time.Now(), flushInterval: time.Duration(math.MaxInt64), periodicFlushContext: context.Background(), now: time.Now}
 }
 
 // NewBufferedLoggerWithFlushInterval returns a new BufferedLogger type by reference. This type will
 // flush logs periodically on Write() calls as well as when Flush*() functions are called.
 func NewBufferedLoggerWithFlushInterval(ctx context.Context, logger *Logger, interval time.Duration) *BufferedLogger {
-	return &BufferedLogger{logger: logger, lastFlush: time.Now(), flushInterval: interval, periodicFlushContext: ctx}
+	return &BufferedLogger{logger: logger, lastFlush: time.Now(), flushInterval: interval, periodicFlushContext: ctx, now: time.Now}
 }
 
 // Write implements the io.Writer interface, converting input to a string
@@ -62,7 +63,7 @@ func (b *BufferedLogger) Write(p []byte) (int, error) {
 	}
 	b.logs = append(b.logs, b.builder.String())
 	b.builder.Reset()
-	if time.Since(b.lastFlush) > b.flushInterval {
+	if b.now().Sub(b.lastFlush) > b.flushInterval {
 		b.FlushAtDebug(b.periodicFlushContext)
 	}
 	return n, err
diff --git a/sdks/go/container/tools/buffered_logging_test.go b/sdks/go/container/tools/buffered_logging_test.go
index 10d5fcf90041..9f542d2d5ab6 100644
--- a/sdks/go/container/tools/buffered_logging_test.go
+++ b/sdks/go/container/tools/buffered_logging_test.go
@@ -189,14 +189,18 @@ func TestBufferedLogger(t *testing.T) {
 	t.Run("debug flush at interval", func(t *testing.T) {
 		catcher := &logCatcher{}
 		l := &Logger{client: catcher}
-		interval := time.Duration(5e9)
+		interval := 5 * time.Second
 		bl := NewBufferedLoggerWithFlushInterval(context.Background(), l, interval)
 
-		messages := []string{"foo", "bar", "baz"}
+		startTime := time.Now()
+		bl.now = func() time.Time { return startTime }
 
-		for _, message := range messages {
-			// Pause for four seconds before each message
-			time.Sleep(2e9)
+		messages := []string{"foo", "bar"}
+
+		for i, message := range messages {
+			if i > 1 {
+				bl.now = func() time.Time { return startTime.Add(6 * time.Second) }
+			}
 			messBytes := []byte(message)
 			n, err := bl.Write(messBytes)
 
@@ -208,9 +212,21 @@ func TestBufferedLogger(t *testing.T) {
 			}
 		}
 
+		lastMessage := "baz"
+		bl.now = func() time.Time { return startTime.Add(6 * time.Second) }
+		messBytes := []byte(lastMessage)
+		n, err := bl.Write(messBytes)
+
+		if err != nil {
+			t.Errorf("got error %v", err)
+		}
+		if got, want := n, len(messBytes); got != want {
+			t.Errorf("got %d bytes written, want %d", got, want)
+		}
+
 		// Type should have auto-flushed at debug after the third message
-		// (12 seconds > 10 second interval)
 		received := catcher.msgs[0].GetLogEntries()
+		messages = append(messages, lastMessage)
 
 		for i, message := range received {
 			if got, want := message.Message, messages[i]; got != want {
diff --git a/sdks/python/container/piputil.go b/sdks/python/container/piputil.go
index 79daa7476e58..045666095982 100644
--- a/sdks/python/container/piputil.go
+++ b/sdks/python/container/piputil.go
@@ -77,7 +77,7 @@ func isPackageInstalled(pkgName string) bool {
 	return true
 }
 
-const pipLogFlushInterval time.Duration = time.Duration(5e9)
+const pipLogFlushInterval time.Duration = 5 * time.Second
 
 // pipInstallPackage installs the given package, if present.
 func pipInstallPackage(ctx context.Context, logger *tools.Logger, files []string, dir, name string, force, optional bool, extras []string) error {