github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/ingestion/parallelizer_test.go (about)

     1  package ingestion_test
     2  
     3  import (
     4  	"context"
     5  	"io/ioutil"
     6  	"sync"
     7  
     8  	. "github.com/onsi/ginkgo/v2"
     9  	. "github.com/onsi/gomega"
    10  	"github.com/pyroscope-io/pyroscope/pkg/ingestion"
    11  	"github.com/pyroscope-io/pyroscope/pkg/storage/metadata"
    12  	"github.com/pyroscope-io/pyroscope/pkg/storage/segment"
    13  	"github.com/pyroscope-io/pyroscope/pkg/util/attime"
    14  	"github.com/sirupsen/logrus"
    15  )
    16  
    17  type mockPutter struct {
    18  	Fn func(context.Context, *ingestion.IngestInput) error
    19  }
    20  
    21  func (m mockPutter) Ingest(ctx context.Context, in *ingestion.IngestInput) error {
    22  	return m.Fn(ctx, in)
    23  }
    24  
    25  var _ = Describe("Parallelizer", func() {
    26  	It("calls Putters", func() {
    27  		logger := logrus.New()
    28  		logger.SetOutput(ioutil.Discard)
    29  
    30  		var wg sync.WaitGroup
    31  
    32  		pi := &ingestion.IngestInput{
    33  			Metadata: ingestion.Metadata{
    34  				Key: segment.NewKey(map[string]string{
    35  					"__name__": "myapp",
    36  				}),
    37  
    38  				StartTime:       attime.Parse("1654110240"),
    39  				EndTime:         attime.Parse("1654110250"),
    40  				SampleRate:      100,
    41  				SpyName:         "gospy",
    42  				Units:           metadata.SamplesUnits,
    43  				AggregationType: metadata.SumAggregationType,
    44  			},
    45  		}
    46  
    47  		fn := func(ctx context.Context, putinput *ingestion.IngestInput) error {
    48  			defer GinkgoRecover()
    49  
    50  			Expect(putinput).To(Equal(pi))
    51  			wg.Done()
    52  			return nil
    53  		}
    54  
    55  		mock1 := mockPutter{Fn: fn}
    56  		mock2 := mockPutter{Fn: fn}
    57  
    58  		p := ingestion.NewParallelizer(logger, mock1, mock2)
    59  
    60  		wg.Add(2)
    61  		p.Ingest(context.TODO(), pi)
    62  		wg.Wait()
    63  	})
    64  })