github.com/mier85/go-sensor@v1.30.1-0.20220920111756-9bf41b3bc7e0/integration_test.go (about)

     1  // (c) Copyright IBM Corp. 2021
     2  // (c) Copyright Instana Inc. 2020
     3  
     4  // +build integration
     5  
     6  package instana_test
     7  
     8  import (
     9  	"fmt"
    10  	"io/ioutil"
    11  	"log"
    12  	"net"
    13  	"net/http"
    14  	"os"
    15  )
    16  
    17  type serverlessAgentPluginPayload struct {
    18  	EntityID string
    19  	Data     map[string]interface{}
    20  }
    21  
    22  type serverlessAgentRequest struct {
    23  	Header http.Header
    24  	Body   []byte
    25  }
    26  
    27  type serverlessAgent struct {
    28  	Bundles []serverlessAgentRequest
    29  
    30  	ln           net.Listener
    31  	restoreEnvFn func()
    32  }
    33  
    34  func setupServerlessAgent() (*serverlessAgent, error) {
    35  	ln, err := net.Listen("tcp", "127.0.0.1:0")
    36  	if err != nil {
    37  		return nil, fmt.Errorf("failed to initialize the serverless agent listener: %s", err)
    38  	}
    39  
    40  	srv := &serverlessAgent{
    41  		ln:           ln,
    42  		restoreEnvFn: restoreEnvVarFunc("INSTANA_ENDPOINT_URL"),
    43  	}
    44  
    45  	mux := http.NewServeMux()
    46  	mux.HandleFunc("/bundle", srv.HandleBundle)
    47  
    48  	go http.Serve(ln, mux)
    49  
    50  	os.Setenv("INSTANA_ENDPOINT_URL", "http://"+ln.Addr().String())
    51  
    52  	return srv, nil
    53  }
    54  
    55  func (srv *serverlessAgent) HandleBundle(w http.ResponseWriter, req *http.Request) {
    56  	body, err := ioutil.ReadAll(req.Body)
    57  	if err != nil {
    58  		log.Printf("ERROR: failed to read serverless agent spans request body: %s", err)
    59  		body = nil
    60  	}
    61  
    62  	srv.Bundles = append(srv.Bundles, serverlessAgentRequest{
    63  		Header: req.Header,
    64  		Body:   body,
    65  	})
    66  
    67  	w.WriteHeader(http.StatusNoContent)
    68  }
    69  
    70  func (srv *serverlessAgent) Reset() {
    71  	srv.Bundles = nil
    72  }
    73  
    74  func (srv *serverlessAgent) Teardown() {
    75  	srv.restoreEnvFn()
    76  	srv.ln.Close()
    77  }
    78  
    79  func restoreEnvVarFunc(key string) func() {
    80  	if oldValue, ok := os.LookupEnv(key); ok {
    81  		return func() { os.Setenv(key, oldValue) }
    82  	}
    83  
    84  	return func() { os.Unsetenv(key) }
    85  }