github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/agent/upstream/remote/remote_test.go (about)

     1  package remote
     2  
     3  import (
     4  	"fmt"
     5  	"html"
     6  	"io"
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"sync"
    10  	"time"
    11  
    12  	. "github.com/onsi/ginkgo/v2"
    13  	. "github.com/onsi/gomega"
    14  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
    15  	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
    16  	"github.com/pyroscope-io/pyroscope/pkg/structs/transporttrie"
    17  	"github.com/pyroscope-io/pyroscope/pkg/testing"
    18  	"github.com/sirupsen/logrus"
    19  )
    20  
    21  var _ = Describe("remote.Remote", func() {
    22  	Describe("Upload", func() {
    23  		It("uploads data to an http server", func() {
    24  			done := make(chan interface{})
    25  			func() {
    26  				defer GinkgoRecover()
    27  
    28  				wg := sync.WaitGroup{}
    29  				wg.Add(3)
    30  				var timestampsMutex sync.Mutex
    31  				timestamps := []time.Time{}
    32  				myHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    33  					defer GinkgoRecover()
    34  
    35  					timestampsMutex.Lock()
    36  					timestamps = append(timestamps, time.Now())
    37  					timestampsMutex.Unlock()
    38  					_, err := io.ReadAll(r.Body)
    39  					Expect(err).ToNot(HaveOccurred())
    40  
    41  					fmt.Fprintf(w, "Hello, %q", html.EscapeString(r.URL.Path))
    42  					wg.Done()
    43  				})
    44  
    45  				httpServer := httptest.NewServer(myHandler)
    46  				defer httpServer.Close()
    47  
    48  				cfg := RemoteConfig{
    49  					AuthToken:              "",
    50  					UpstreamThreads:        4,
    51  					UpstreamAddress:        httpServer.URL,
    52  					UpstreamRequestTimeout: 3 * time.Second,
    53  				}
    54  				r, err := New(cfg, logrus.New())
    55  				r.Start()
    56  
    57  				t := transporttrie.New()
    58  				for i := 0; i < 3; i++ {
    59  					r.Upload(&upstream.UploadJob{
    60  						Name:       "test{}",
    61  						StartTime:  testing.SimpleTime(0),
    62  						EndTime:    testing.SimpleTime(10),
    63  						SpyName:    "debugspy",
    64  						SampleRate: 100,
    65  						Units:      metadata.SamplesUnits,
    66  						Trie:       t,
    67  					})
    68  				}
    69  
    70  				Expect(err).To(BeNil())
    71  				wg.Wait()
    72  				r.Stop()
    73  				close(done)
    74  			}()
    75  			Eventually(done, 5).Should(BeClosed())
    76  		})
    77  	})
    78  })