github.com/newrelic/go-agent@v3.26.0+incompatible/internal_response_writer_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  	"bufio"
     8  	"io"
     9  	"net"
    10  	"net/http"
    11  	"testing"
    12  )
    13  
    14  type rwNoExtraMethods struct {
    15  	hijackCalled      bool
    16  	readFromCalled    bool
    17  	flushCalled       bool
    18  	closeNotifyCalled bool
    19  }
    20  
    21  type rwTwoExtraMethods struct{ rwNoExtraMethods }
    22  type rwAllExtraMethods struct{ rwTwoExtraMethods }
    23  
    24  func (rw *rwAllExtraMethods) CloseNotify() <-chan bool {
    25  	rw.closeNotifyCalled = true
    26  	return nil
    27  }
    28  func (rw *rwAllExtraMethods) ReadFrom(r io.Reader) (int64, error) {
    29  	rw.readFromCalled = true
    30  	return 0, nil
    31  }
    32  
    33  func (rw *rwNoExtraMethods) Header() http.Header        { return nil }
    34  func (rw *rwNoExtraMethods) Write([]byte) (int, error)  { return 0, nil }
    35  func (rw *rwNoExtraMethods) WriteHeader(statusCode int) {}
    36  
    37  func (rw *rwTwoExtraMethods) Flush() {
    38  	rw.flushCalled = true
    39  }
    40  func (rw *rwTwoExtraMethods) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    41  	rw.hijackCalled = true
    42  	return nil, nil, nil
    43  }
    44  
    45  func TestTransactionAllExtraMethods(t *testing.T) {
    46  	app := testApp(nil, nil, t)
    47  	rw := &rwAllExtraMethods{}
    48  	txn := app.StartTransaction("hello", rw, nil)
    49  	if v, ok := txn.(http.CloseNotifier); ok {
    50  		v.CloseNotify()
    51  	}
    52  	if v, ok := txn.(http.Flusher); ok {
    53  		v.Flush()
    54  	}
    55  	if v, ok := txn.(http.Hijacker); ok {
    56  		v.Hijack()
    57  	}
    58  	if v, ok := txn.(io.ReaderFrom); ok {
    59  		v.ReadFrom(nil)
    60  	}
    61  	if !rw.hijackCalled ||
    62  		!rw.readFromCalled ||
    63  		!rw.flushCalled ||
    64  		!rw.closeNotifyCalled {
    65  		t.Error("wrong methods called", rw)
    66  	}
    67  }
    68  
    69  func TestTransactionNoExtraMethods(t *testing.T) {
    70  	app := testApp(nil, nil, t)
    71  	rw := &rwNoExtraMethods{}
    72  	txn := app.StartTransaction("hello", rw, nil)
    73  	if _, ok := txn.(http.CloseNotifier); ok {
    74  		t.Error("unexpected CloseNotifier method")
    75  	}
    76  	if _, ok := txn.(http.Flusher); ok {
    77  		t.Error("unexpected Flusher method")
    78  	}
    79  	if _, ok := txn.(http.Hijacker); ok {
    80  		t.Error("unexpected Hijacker method")
    81  	}
    82  	if _, ok := txn.(io.ReaderFrom); ok {
    83  		t.Error("unexpected ReaderFrom method")
    84  	}
    85  }
    86  
    87  func TestTransactionTwoExtraMethods(t *testing.T) {
    88  	app := testApp(nil, nil, t)
    89  	rw := &rwTwoExtraMethods{}
    90  	txn := app.StartTransaction("hello", rw, nil)
    91  	if _, ok := txn.(http.CloseNotifier); ok {
    92  		t.Error("unexpected CloseNotifier method")
    93  	}
    94  	if v, ok := txn.(http.Flusher); ok {
    95  		v.Flush()
    96  	}
    97  	if v, ok := txn.(http.Hijacker); ok {
    98  		v.Hijack()
    99  	}
   100  	if _, ok := txn.(io.ReaderFrom); ok {
   101  		t.Error("unexpected ReaderFrom method")
   102  	}
   103  	if !rw.hijackCalled ||
   104  		rw.readFromCalled ||
   105  		!rw.flushCalled ||
   106  		rw.closeNotifyCalled {
   107  		t.Error("wrong methods called", rw)
   108  	}
   109  }