Repository: prometheus/prometheus
File: storage/remote/queue_manager_test.go (Go)

Code before:
		tsID := getSeriesIDFromRef(series[h.Ref])
		c.expectedHistograms[tsID] = append(c.expectedHistograms[tsID], prompb.FromIntHistogram(h.T, h.H))
	}
}

func (c *TestWriteClient) expectFloatHistograms(fhs []record.RefFloatHistogramSample, series []record.RefSeries) {
	c.mtx.Lock()
	defer c.mtx.Unlock()

	c.expectedFloatHistograms = map[string][]prompb.Histogram{}
	c.receivedFloatHistograms = map[string][]prompb.Histogram{}

	for _, fh := range fhs {
		tsID := getSeriesIDFromRef(series[fh.Ref])
		c.expectedFloatHistograms[tsID] = append(c.expectedFloatHistograms[tsID], prompb.FromFloatHistogram(fh.T, fh.FH))
	}
}

func (c *TestWriteClient) expectMetadataForBatch(metadata []record.RefMetadata, series []record.RefSeries, samples []record.RefSample, exemplars []record.RefExemplar, histograms []record.RefHistogramSample, floatHistograms []record.RefFloatHistogramSample) {
	c.mtx.Lock()
	defer c.mtx.Unlock()

	c.expectedMetadata = map[string][]prompb.MetricMetadata{}
	c.receivedMetadata = map[string][]prompb.MetricMetadata{}

	// Collect refs that have data in this batch
	refsWithData := make(map[chunks.HeadSeriesRef]bool)
	for _, s := range samples {
		refsWithData[s.Ref] = true
	}
	for _, e := range exemplars {
		refsWithData[e.Ref] = true
	}
	for _, h := range histograms {
		refsWithData[h.Ref] = true
	}
	for _, fh := range floatHistograms {
		refsWithData[fh.Ref] = true
	}

	// Only expect metadata for series that have data in this batch
	for _, m := range metadata {
		if !refsWithData[m.Ref] {
			continue
		}
		tsID := getSeriesIDFromRef(series[m.Ref])
		c.expectedMetadata[tsID] = append(c.expectedMetadata[tsID], prompb.MetricMetadata{
			MetricFamilyName: tsID,
			Type:             prompb.FromMetadataType(record.ToMetricType(m.Type)),
			Help:             m.Help,
			Unit:             m.Unit,

Review comment (performance):
```suggestion
	// Collect refs that have data in this batch.
```

Code after:
		tsID := getSeriesIDFromRef(series[h.Ref])
		c.expectedHistograms[tsID] = append(c.expectedHistograms[tsID], prompb.FromIntHistogram(h.T, h.H))
	}
}

func (c *TestWriteClient) expectFloatHistograms(fhs []record.RefFloatHistogramSample, series []record.RefSeries) {
	c.mtx.Lock()
	defer c.mtx.Unlock()

	c.expectedFloatHistograms = map[string][]prompb.Histogram{}
	c.receivedFloatHistograms = map[string][]prompb.Histogram{}

	for _, fh := range fhs {
		tsID := getSeriesIDFromRef(series[fh.Ref])
		c.expectedFloatHistograms[tsID] = append(c.expectedFloatHistograms[tsID], prompb.FromFloatHistogram(fh.T, fh.FH))
	}
}

func (c *TestWriteClient) expectMetadataForBatch(metadata []record.RefMetadata, series []record.RefSeries, samples []record.RefSample, exemplars []record.RefExemplar, histograms []record.RefHistogramSample, floatHistograms []record.RefFloatHistogramSample) {
	c.mtx.Lock()
	defer c.mtx.Unlock()

	c.expectedMetadata = map[string][]prompb.MetricMetadata{}
	c.receivedMetadata = map[string][]prompb.MetricMetadata{}

	// Collect refs that have data in this batch.
	refsWithData := make(map[chunks.HeadSeriesRef]struct{})
	for _, s := range samples {
		refsWithData[s.Ref] = struct{}{}
	}
	for _, e := range exemplars {
		refsWithData[e.Ref] = struct{}{}
	}
	for _, h := range histograms {
		refsWithData[h.Ref] = struct{}{}
	}
	for _, fh := range floatHistograms {
		refsWithData[fh.Ref] = struct{}{}
	}

	// Only expect metadata for series that have data in this batch.
	for _, m := range metadata {
		if _, ok := refsWithData[m.Ref]; !ok {
			continue
		}
		tsID := getSeriesIDFromRef(series[m.Ref])
		c.expectedMetadata[tsID] = append(c.expectedMetadata[tsID], prompb.MetricMetadata{
			MetricFamilyName: tsID,
			Type:             prompb.FromMetadataType(record.ToMetricType(m.Type)),
			Help:             m.Help,
			Unit:             m.Unit,