88 "context"
99 "encoding/json"
1010 "errors"
11+ "io"
1112 "net/http"
13+ "os"
1214 "strings"
1315 "testing"
1416
@@ -45,14 +47,23 @@ func distributedTracingEnabled(key string) string {
4547 }
4648}
4749
50+ // bufWriterProvider is a testing implementation of writerProvider
51+ type bufWriterProvider struct {
52+ buf io.Writer
53+ }
54+
55+ func (bw bufWriterProvider ) borrowWriter (needsWriter func (writer io.Writer )) {
56+ needsWriter (bw .buf )
57+ }
58+
4859func TestColdStart (t * testing.T ) {
4960 originalHandler := func (c context.Context ) {}
5061 app := testApp (nil , t )
5162 wrapped := Wrap (originalHandler , app )
5263 w := wrapped .(* wrappedHandler )
5364 w .functionName = "functionName"
5465 buf := & bytes.Buffer {}
55- w .writer = buf
66+ w .hasWriter = bufWriterProvider { buf }
5667
5768 ctx := context .Background ()
5869 lctx := & lambdacontext.LambdaContext {
@@ -104,7 +115,7 @@ func TestColdStart(t *testing.T) {
104115
105116 // Invoke the handler again to test the cold-start attribute absence.
106117 buf = & bytes.Buffer {}
107- w .writer = buf
118+ w .hasWriter = bufWriterProvider { buf }
108119 internal .HarvestTesting (app .Private , nil )
109120 resp , err = wrapped .Invoke (ctx , nil )
110121 if nil != err || string (resp ) != "null" {
@@ -154,7 +165,7 @@ func TestErrorCapture(t *testing.T) {
154165 w := wrapped .(* wrappedHandler )
155166 w .functionName = "functionName"
156167 buf := & bytes.Buffer {}
157- w .writer = buf
168+ w .hasWriter = bufWriterProvider { buf }
158169
159170 resp , err := wrapped .Invoke (context .Background (), nil )
160171 if err != returnError || string (resp ) != "" {
@@ -229,7 +240,7 @@ func TestSetWebRequest(t *testing.T) {
229240 w := wrapped .(* wrappedHandler )
230241 w .functionName = "functionName"
231242 buf := & bytes.Buffer {}
232- w .writer = buf
243+ w .hasWriter = bufWriterProvider { buf }
233244
234245 req := events.APIGatewayProxyRequest {
235246 Headers : map [string ]string {
@@ -301,7 +312,7 @@ func TestDistributedTracing(t *testing.T) {
301312 w := wrapped .(* wrappedHandler )
302313 w .functionName = "functionName"
303314 buf := & bytes.Buffer {}
304- w .writer = buf
315+ w .hasWriter = bufWriterProvider { buf }
305316
306317 dtHdr := http.Header {}
307318 app .StartTransaction ("hello" ).InsertDistributedTraceHeaders (dtHdr )
@@ -396,7 +407,7 @@ func TestEventARN(t *testing.T) {
396407 w := wrapped .(* wrappedHandler )
397408 w .functionName = "functionName"
398409 buf := & bytes.Buffer {}
399- w .writer = buf
410+ w .hasWriter = bufWriterProvider { buf }
400411
401412 req := events.DynamoDBEvent {
402413 Records : []events.DynamoDBEventRecord {{
@@ -473,7 +484,7 @@ func TestAPIGatewayProxyResponse(t *testing.T) {
473484 w := wrapped .(* wrappedHandler )
474485 w .functionName = "functionName"
475486 buf := & bytes.Buffer {}
476- w .writer = buf
487+ w .hasWriter = bufWriterProvider { buf }
477488
478489 resp , err := wrapped .Invoke (context .Background (), nil )
479490 if nil != err {
@@ -535,7 +546,7 @@ func TestCustomEvent(t *testing.T) {
535546 w := wrapped .(* wrappedHandler )
536547 w .functionName = "functionName"
537548 buf := & bytes.Buffer {}
538- w .writer = buf
549+ w .hasWriter = bufWriterProvider { buf }
539550
540551 resp , err := wrapped .Invoke (context .Background (), nil )
541552 if nil != err || string (resp ) != "null" {
@@ -555,3 +566,30 @@ func TestCustomEvent(t *testing.T) {
555566 t .Error ("no output written" )
556567 }
557568}
569+
570+ func TestDefaultWriterProvider (t * testing.T ) {
571+ dwp := defaultWriterProvider {}
572+ dwp .borrowWriter (func (writer io.Writer ) {
573+ if writer != os .Stdout {
574+ t .Error ("Expected stdout" )
575+ }
576+ })
577+
578+ const telemetryFile = "/tmp/newrelic-telemetry"
579+ defer os .Remove (telemetryFile )
580+ file , err := os .Create (telemetryFile )
581+ if err != nil {
582+ t .Error ("Unexpected error creating telemetry file" , err )
583+ }
584+
585+ err = file .Close ()
586+ if err != nil {
587+ t .Error ("Error closing telemetry file" , err )
588+ }
589+
590+ dwp .borrowWriter (func (writer io.Writer ) {
591+ if writer == os .Stdout {
592+ t .Error ("Expected telemetry file, got stdout" )
593+ }
594+ })
595+ }
0 commit comments