github.com/newrelic/go-agent@v3.26.0+incompatible/internal_set_web_response_test.go (about)

     1  // Copyright 2020 New Relic Corporation. All rights reserved.
     2  // SPDX-License-Identifier: Apache-2.0
     3  
     4  package newrelic
     5  
     6  import (
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"testing"
    10  
    11  	"github.com/newrelic/go-agent/internal"
    12  )
    13  
    14  func TestTransactionStartedWithoutResponse(t *testing.T) {
    15  	// Test that the http.ResponseWriter methods of the transaction can
    16  	// safely be called if a ResponseWriter is not provided.
    17  	app := testApp(nil, nil, t)
    18  	txn := app.StartTransaction("hello", nil, nil)
    19  	txn.WriteHeader(123)
    20  	if hdr := txn.Header(); hdr != nil {
    21  		t.Error(hdr)
    22  	}
    23  	n, err := txn.Write([]byte("should not panic"))
    24  	if err != nil || n != 0 {
    25  		t.Error(err, n)
    26  	}
    27  	txn.End()
    28  	app.ExpectTxnEvents(t, []internal.WantEvent{{
    29  		AgentAttributes: map[string]interface{}{"httpResponseCode": 123},
    30  		Intrinsics:      map[string]interface{}{"name": "OtherTransaction/Go/hello"},
    31  	}})
    32  }
    33  
    34  func TestSetWebResponseNil(t *testing.T) {
    35  	// Test that the http.ResponseWriter methods of the transaction can
    36  	// safely be called if txn.SetWebResponse(nil) has been called.
    37  	app := testApp(nil, nil, t)
    38  	txn := app.StartTransaction("hello", nil, nil)
    39  	txn = txn.SetWebResponse(nil)
    40  	txn.WriteHeader(123)
    41  	if hdr := txn.Header(); hdr != nil {
    42  		t.Error(hdr)
    43  	}
    44  	n, err := txn.Write([]byte("should not panic"))
    45  	if err != nil || n != 0 {
    46  		t.Error(err, n)
    47  	}
    48  	txn.End()
    49  	app.ExpectTxnEvents(t, []internal.WantEvent{{
    50  		AgentAttributes: map[string]interface{}{"httpResponseCode": 123},
    51  		Intrinsics:      map[string]interface{}{"name": "OtherTransaction/Go/hello"},
    52  	}})
    53  }
    54  
    55  func TestSetWebResponseSuccess(t *testing.T) {
    56  	// Test that the http.ResponseWriter methods of the transaction use the
    57  	// writer set by SetWebResponse.
    58  	app := testApp(nil, nil, t)
    59  	txn := app.StartTransaction("hello", nil, nil)
    60  	w := httptest.NewRecorder()
    61  	txn = txn.SetWebResponse(w)
    62  	txn.WriteHeader(123)
    63  	hdr := txn.Header()
    64  	hdr.Set("zip", "zap")
    65  	body := "should not panic"
    66  	n, err := txn.Write([]byte(body))
    67  	if err != nil || n != len(body) {
    68  		t.Error(err, n)
    69  	}
    70  	txn.End()
    71  	if w.Code != 123 {
    72  		t.Error(w.Code)
    73  	}
    74  	if w.HeaderMap.Get("zip") != "zap" {
    75  		t.Error(w.HeaderMap)
    76  	}
    77  	if w.Body.String() != body {
    78  		t.Error(w.Body.String())
    79  	}
    80  	app.ExpectTxnEvents(t, []internal.WantEvent{{
    81  		AgentAttributes: map[string]interface{}{"httpResponseCode": 123},
    82  		Intrinsics:      map[string]interface{}{"name": "OtherTransaction/Go/hello"},
    83  	}})
    84  }
    85  
    86  type writerWithFlush struct{}
    87  
    88  func (w writerWithFlush) Header() http.Header       { return nil }
    89  func (w writerWithFlush) WriteHeader(int)           {}
    90  func (w writerWithFlush) Write([]byte) (int, error) { return 0, nil }
    91  func (w writerWithFlush) Flush()                    {}
    92  
    93  func TestSetWebResponseTxnUpgraded(t *testing.T) {
    94  	// Test that the using Transaction reference returned by SetWebResponse
    95  	// properly has the optional methods that the ResponseWriter does.
    96  	app := testApp(nil, nil, t)
    97  	txn := app.StartTransaction("hello", nil, nil)
    98  	if _, ok := txn.(http.Flusher); ok {
    99  		t.Error("should not have Flusher")
   100  	}
   101  	txn = txn.SetWebResponse(writerWithFlush{})
   102  	if _, ok := txn.(http.Flusher); !ok {
   103  		t.Error("should have Flusher now")
   104  	}
   105  }