github.com/flavio/docker@v0.1.3-0.20170117145210-f63d1a6eec47/daemon/logger/splunk/splunkhecmock_test.go (about)

     1  package splunk
     2  
     3  import (
     4  	"compress/gzip"
     5  	"encoding/json"
     6  	"fmt"
     7  	"io"
     8  	"io/ioutil"
     9  	"net"
    10  	"net/http"
    11  	"testing"
    12  )
    13  
    14  func (message *splunkMessage) EventAsString() (string, error) {
    15  	if val, ok := message.Event.(string); ok {
    16  		return val, nil
    17  	}
    18  	return "", fmt.Errorf("Cannot cast Event %v to string", message.Event)
    19  }
    20  
    21  func (message *splunkMessage) EventAsMap() (map[string]interface{}, error) {
    22  	if val, ok := message.Event.(map[string]interface{}); ok {
    23  		return val, nil
    24  	}
    25  	return nil, fmt.Errorf("Cannot cast Event %v to map", message.Event)
    26  }
    27  
    28  type HTTPEventCollectorMock struct {
    29  	tcpAddr     *net.TCPAddr
    30  	tcpListener *net.TCPListener
    31  
    32  	token               string
    33  	simulateServerError bool
    34  
    35  	test *testing.T
    36  
    37  	connectionVerified bool
    38  	gzipEnabled        *bool
    39  	messages           []*splunkMessage
    40  	numOfRequests      int
    41  }
    42  
    43  func NewHTTPEventCollectorMock(t *testing.T) *HTTPEventCollectorMock {
    44  	tcpAddr := &net.TCPAddr{IP: []byte{127, 0, 0, 1}, Port: 0, Zone: ""}
    45  	tcpListener, err := net.ListenTCP("tcp", tcpAddr)
    46  	if err != nil {
    47  		t.Fatal(err)
    48  	}
    49  	return &HTTPEventCollectorMock{
    50  		tcpAddr:             tcpAddr,
    51  		tcpListener:         tcpListener,
    52  		token:               "4642492F-D8BD-47F1-A005-0C08AE4657DF",
    53  		simulateServerError: false,
    54  		test:                t,
    55  		connectionVerified:  false}
    56  }
    57  
    58  func (hec *HTTPEventCollectorMock) URL() string {
    59  	return "http://" + hec.tcpListener.Addr().String()
    60  }
    61  
    62  func (hec *HTTPEventCollectorMock) Serve() error {
    63  	return http.Serve(hec.tcpListener, hec)
    64  }
    65  
    66  func (hec *HTTPEventCollectorMock) Close() error {
    67  	return hec.tcpListener.Close()
    68  }
    69  
    70  func (hec *HTTPEventCollectorMock) ServeHTTP(writer http.ResponseWriter, request *http.Request) {
    71  	var err error
    72  
    73  	hec.numOfRequests++
    74  
    75  	if hec.simulateServerError {
    76  		if request.Body != nil {
    77  			defer request.Body.Close()
    78  		}
    79  		writer.WriteHeader(http.StatusInternalServerError)
    80  		return
    81  	}
    82  
    83  	switch request.Method {
    84  	case http.MethodOptions:
    85  		// Verify that options method is getting called only once
    86  		if hec.connectionVerified {
    87  			hec.test.Errorf("Connection should not be verified more than once. Got second request with %s method.", request.Method)
    88  		}
    89  		hec.connectionVerified = true
    90  		writer.WriteHeader(http.StatusOK)
    91  	case http.MethodPost:
    92  		// Always verify that Driver is using correct path to HEC
    93  		if request.URL.String() != "/services/collector/event/1.0" {
    94  			hec.test.Errorf("Unexpected path %v", request.URL)
    95  		}
    96  		defer request.Body.Close()
    97  
    98  		if authorization, ok := request.Header["Authorization"]; !ok || authorization[0] != ("Splunk "+hec.token) {
    99  			hec.test.Error("Authorization header is invalid.")
   100  		}
   101  
   102  		gzipEnabled := false
   103  		if contentEncoding, ok := request.Header["Content-Encoding"]; ok && contentEncoding[0] == "gzip" {
   104  			gzipEnabled = true
   105  		}
   106  
   107  		if hec.gzipEnabled == nil {
   108  			hec.gzipEnabled = &gzipEnabled
   109  		} else if gzipEnabled != *hec.gzipEnabled {
   110  			// Nothing wrong with that, but we just know that Splunk Logging Driver does not do that
   111  			hec.test.Error("Driver should not change Content Encoding.")
   112  		}
   113  
   114  		var gzipReader *gzip.Reader
   115  		var reader io.Reader
   116  		if gzipEnabled {
   117  			gzipReader, err = gzip.NewReader(request.Body)
   118  			if err != nil {
   119  				hec.test.Fatal(err)
   120  			}
   121  			reader = gzipReader
   122  		} else {
   123  			reader = request.Body
   124  		}
   125  
   126  		// Read body
   127  		var body []byte
   128  		body, err = ioutil.ReadAll(reader)
   129  		if err != nil {
   130  			hec.test.Fatal(err)
   131  		}
   132  
   133  		// Parse message
   134  		messageStart := 0
   135  		for i := 0; i < len(body); i++ {
   136  			if i == len(body)-1 || (body[i] == '}' && body[i+1] == '{') {
   137  				var message splunkMessage
   138  				err = json.Unmarshal(body[messageStart:i+1], &message)
   139  				if err != nil {
   140  					hec.test.Log(string(body[messageStart : i+1]))
   141  					hec.test.Fatal(err)
   142  				}
   143  				hec.messages = append(hec.messages, &message)
   144  				messageStart = i + 1
   145  			}
   146  		}
   147  
   148  		if gzipEnabled {
   149  			gzipReader.Close()
   150  		}
   151  
   152  		writer.WriteHeader(http.StatusOK)
   153  	default:
   154  		hec.test.Errorf("Unexpected HTTP method %s", http.MethodOptions)
   155  		writer.WriteHeader(http.StatusBadRequest)
   156  	}
   157  }