fix: race-condition problem on testhelper (glassonion1#5)

internal/config/config.go

@@ -1,6 +1,7 @@
 package config
 
 import (
+	"context"
 	"io"
 	"os"
 
@@ -28,3 +29,25 @@ func init() {
 	ApplicationLogOut = os.Stdout
 	AccessLogOut = os.Stderr
 }
+
+type ContextConfig struct {
+	// ApplicationLogOut is io.Writer object for application log
+	ApplicationLogOut io.Writer
+	// AccessLogOut is io.Writer object for access log
+	AccessLogOut io.Writer
+}
+
+type contextKey struct{}
+
+var contextConfigKey = &contextKey{}
+
+// GetContextConfig sets the ContextConfig instance to context
+func SetContextConfig(ctx context.Context, cs *ContextConfig) context.Context {
+	return context.WithValue(ctx, contextConfigKey, cs)
+}
+
+// GetContextConfig gets the ContextSeverity instance from context
+func GetContextConfig(ctx context.Context) *ContextConfig {
+	v, _ := ctx.Value(contextConfigKey).(*ContextConfig)
+	return v
+}

internal/logger/logger.go

@@ -53,7 +53,12 @@ func WriteApplicationLog(ctx context.Context, s severity.Severity, format string
 		TraceSampled:   sc.TraceSampled,
 	}
 
-	if err := json.NewEncoder(config.ApplicationLogOut).Encode(ety); err != nil {
+	w := config.ApplicationLogOut
+	if cc := config.GetContextConfig(ctx); cc != nil && cc.ApplicationLogOut != nil {
+		w = cc.ApplicationLogOut
+	}
+
+	if err := json.NewEncoder(w).Encode(ety); err != nil {
 		fmt.Printf("failed to write log: %v", err)
 	}
 }
@@ -86,7 +91,12 @@ func WriteAccessLog(ctx context.Context, req types.HTTPRequest) {
 		HTTPRequest: req,
 	}
 
-	if err := json.NewEncoder(config.AccessLogOut).Encode(ety); err != nil {
+	w := config.AccessLogOut
+	if cc := config.GetContextConfig(ctx); cc != nil && cc.AccessLogOut != nil {
+		w = cc.AccessLogOut
+	}
+
+	if err := json.NewEncoder(w).Encode(ety); err != nil {
 		fmt.Printf("failed to write log: %v", err)
 	}
 }

internal/logger/logger_test.go

@@ -48,10 +48,10 @@ func TestLoggerWriteApplicationLog(t *testing.T) {
 	}()
 
 	t.Run("Tests WriteApplicationLog function", func(t *testing.T) {
-		got := testhelper.ExtractApplicationLogOut(t, func() {
-			// tests the function
-			logger.WriteApplicationLog(ctx, severity.Info, "writes %s log", "info")
-		})
+		got := testhelper.ExtractApplicationLogOut(t, ctx, func(ctx context.Context) {
+				// tests the function
+				logger.WriteApplicationLog(ctx, severity.Info, "writes %s log", "info")
+			})
 
 		expected := `{"severity":"INFO","message":"writes info log","time":"2020-12-31T23:59:59.999999999Z","logging.googleapis.com/sourceLocation":{"file":"logger_test.go","line":"51","function":"github.com/glassonion1/logz/internal/logger_test.TestLoggerWriteApplicationLog.func3"},"logging.googleapis.com/trace":"projects/test/traces/00000000000000000000000000000000","logging.googleapis.com/spanId":"0000000000000000","logging.googleapis.com/trace_sampled":false}`
 
@@ -99,7 +99,7 @@ func TestLoggerWriteAccessLog(t *testing.T) {
 
 	t.Run("Tests WriteAccessLog function", func(t *testing.T) {
 
-		got := testhelper.ExtractAccessLogOut(t, func() {
+		got := testhelper.ExtractAccessLogOut(t, ctx, func(ctx context.Context) {
 			// Tests the function
 			httpReq := httptest.NewRequest(http.MethodGet, "/test1", nil)
 			req := types.MakeHTTPRequest(*httpReq, 200, 333, time.Duration(100))
@@ -112,4 +112,5 @@ func TestLoggerWriteAccessLog(t *testing.T) {
 			t.Errorf("failed log info test: %v", diff)
 		}
 	})
+
 }

logz_test.go

@@ -130,7 +130,7 @@ func TestLogzWriteLog(t *testing.T) {
 
 			sc := spancontext.Extract(ctx)
 
-			got := testhelper.ExtractApplicationLogOut(t, func() {
+			got := testhelper.ExtractApplicationLogOut(t, ctx, func(ctx context.Context) {
 				logz.Infof(ctx, "writes %s log", "info")
 			})
 

middleware/nethttp_test.go

@@ -1,6 +1,7 @@
 package middleware_test
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -65,13 +66,14 @@ func TestNetHTTPMaxSeverity(t *testing.T) {
 		mid := middleware.NetHTTP("test/component")(mux)
 		rec := httptest.NewRecorder()
 
-		got := testhelper.ExtractAccessLogOut(t, func() {
+
+		got := testhelper.ExtractAccessLogOut(t, context.Background(), func(ctx context.Context) {
 			req1 := httptest.NewRequest(http.MethodGet, "/test1", nil)
-			mid.ServeHTTP(rec, req1)
+			mid.ServeHTTP(rec, req1.WithContext(ctx))
 			req2 := httptest.NewRequest(http.MethodGet, "/test2", nil)
-			mid.ServeHTTP(rec, req2)
+			mid.ServeHTTP(rec, req2.WithContext(ctx))
 			req3 := httptest.NewRequest(http.MethodGet, "/test3", nil)
-			mid.ServeHTTP(rec, req3)
+			mid.ServeHTTP(rec, req3.WithContext(ctx))
 		})
 
 		if !strings.Contains(got, `"severity":"ERROR"`) {
@@ -107,13 +109,13 @@ func TestNetHTTPMaxSeverityNoLog(t *testing.T) {
 		mid := middleware.NetHTTP("test/component")(mux)
 		rec := httptest.NewRecorder()
 
-		got := testhelper.ExtractAccessLogOut(t, func() {
+		got := testhelper.ExtractAccessLogOut(t, context.Background(), func(ctx context.Context) {
 			req1 := httptest.NewRequest(http.MethodGet, "/test1", nil)
-			mid.ServeHTTP(rec, req1)
+			mid.ServeHTTP(rec, req1.WithContext(ctx))
 			req2 := httptest.NewRequest(http.MethodGet, "/test2", nil)
-			mid.ServeHTTP(rec, req2)
+			mid.ServeHTTP(rec, req2.WithContext(ctx))
 			req3 := httptest.NewRequest(http.MethodGet, "/test3", nil)
-			mid.ServeHTTP(rec, req3)
+			mid.ServeHTTP(rec, req3.WithContext(ctx))
 		})
 
 		if !strings.Contains(got, `"severity":"ERROR"`) {

testhelper/testhelper.go

@@ -2,39 +2,51 @@ package testhelper
 
 import (
 	"bytes"
-	"os"
+	"context"
+	"io"
 	"strings"
 	"testing"
 
 	"github.com/glassonion1/logz/internal/config"
 )
 
 // ExtractStdout extracts string from stdout
-func ExtractApplicationLogOut(t *testing.T, fnc func()) string {
+func ExtractApplicationLogOut(t *testing.T, ctx context.Context, fnc func(ctx context.Context)) string {
 	t.Helper()
 
 	var buf bytes.Buffer
-	config.ApplicationLogOut = &buf
-	defer func() {
-		config.ApplicationLogOut = os.Stdout
-	}()
-
-	fnc()
+	if cc := config.GetContextConfig(ctx); cc != nil {
+		cc.ApplicationLogOut = &buf
+	} else {
+		ctx = config.SetContextConfig(ctx, &config.ContextConfig{ApplicationLogOut: &buf})
+	}
+	fnc(ctx)
 
 	return strings.TrimRight(buf.String(), "\n")
 }
 
 // ExtractStdout extracts string from stderr
-func ExtractAccessLogOut(t *testing.T, fnc func()) string {
+func ExtractAccessLogOut(t *testing.T, ctx context.Context, fnc func(ctx context.Context)) string {
 	t.Helper()
 
 	var buf bytes.Buffer
-	config.AccessLogOut = &buf
-	defer func() {
-		config.AccessLogOut = os.Stdout
-	}()
+	if cc := config.GetContextConfig(ctx); cc != nil {
+		cc.AccessLogOut = &buf
+	} else {
+		ctx = config.SetContextConfig(ctx, &config.ContextConfig{AccessLogOut: &buf})
+	}
 
-	fnc()
+	fnc(ctx)
 
 	return strings.TrimRight(buf.String(), "\n")
 }
+
+
+// OverrideLogOutContext override log I/O in the context
+func OverrideLogOutContext(t *testing.T, ctx context.Context, appLogOut, accessLogOut io.Writer) context.Context {
+	t.Helper()
+	return config.SetContextConfig(ctx, &config.ContextConfig{
+		ApplicationLogOut: appLogOut,
+		AccessLogOut: accessLogOut,
+	})
+}