github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/block/testutil/mock_block.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/storage/tsdb/testutil/block_mock.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package testutil
     7  
     8  import (
     9  	"context"
    10  	"crypto/rand"
    11  	"encoding/json"
    12  	"fmt"
    13  	"strings"
    14  	"testing"
    15  	"time"
    16  
    17  	"github.com/oklog/ulid/v2"
    18  	"github.com/prometheus/common/model"
    19  	"github.com/stretchr/testify/require"
    20  
    21  	"github.com/grafana/pyroscope/pkg/objstore"
    22  	"github.com/grafana/pyroscope/pkg/phlaredb/block"
    23  )
    24  
    25  func MockStorageBlock(t testing.TB, bucket objstore.Bucket, userID string, minT, maxT model.Time) block.Meta {
    26  	return MockStorageBlockWithExtLabels(t, bucket, userID, minT, maxT, nil)
    27  }
    28  
    29  func MockStorageBlockWithExtLabels(t testing.TB, bucket objstore.Bucket, userID string, minT, maxT model.Time, externalLabels map[string]string) block.Meta {
    30  	// Generate a block ID whose timestamp matches the maxT (for simplicity we assume it
    31  	// has been compacted and shipped in zero time, even if not realistic).
    32  	id := ulid.MustNew(uint64(maxT), rand.Reader)
    33  
    34  	meta := block.Meta{
    35  		Version: 1,
    36  		ULID:    id,
    37  		MinTime: minT,
    38  		MaxTime: maxT,
    39  		Compaction: block.BlockMetaCompaction{
    40  			Level:   1,
    41  			Sources: []ulid.ULID{id},
    42  		},
    43  		Labels: externalLabels,
    44  	}
    45  
    46  	metaContent, err := json.Marshal(meta)
    47  	require.NoError(t, err, "failed to marshal mocked block meta")
    48  
    49  	metaContentReader := strings.NewReader(string(metaContent))
    50  	metaPath := fmt.Sprintf("%s/phlaredb/%s/meta.json", userID, id.String())
    51  	require.NoError(t, bucket.Upload(context.Background(), metaPath, metaContentReader))
    52  
    53  	// Upload an empty index, just to make sure the meta.json is not the only object in the block location.
    54  	indexPath := fmt.Sprintf("%s/phlaredb/%s/index", userID, id.String())
    55  	require.NoError(t, bucket.Upload(context.Background(), indexPath, strings.NewReader("")))
    56  
    57  	return meta
    58  }
    59  
    60  func MockStorageDeletionMark(t testing.TB, bucket objstore.Bucket, userID string, meta block.Meta) *block.DeletionMark {
    61  	mark := block.DeletionMark{
    62  		ID:           meta.ULID,
    63  		DeletionTime: time.Now().Add(-time.Minute).Unix(),
    64  		Version:      block.DeletionMarkVersion1,
    65  	}
    66  
    67  	markContent, err := json.Marshal(mark)
    68  	require.NoError(t, err, "failed to marshal mocked deletion mark")
    69  
    70  	markContentReader := strings.NewReader(string(markContent))
    71  	markPath := fmt.Sprintf("%s/phlaredb/%s/%s", userID, meta.ULID.String(), block.DeletionMarkFilename)
    72  	require.NoError(t, bucket.Upload(context.Background(), markPath, markContentReader))
    73  
    74  	return &mark
    75  }
    76  
    77  func MockNoCompactMark(t testing.TB, bucket objstore.Bucket, userID string, meta block.Meta) *block.NoCompactMark {
    78  	mark := block.NoCompactMark{
    79  		ID:            meta.ULID,
    80  		NoCompactTime: time.Now().Unix(),
    81  		Version:       block.DeletionMarkVersion1,
    82  		Details:       "details",
    83  		Reason:        block.ManualNoCompactReason,
    84  	}
    85  
    86  	markContent, err := json.Marshal(mark)
    87  	require.NoError(t, err, "failed to marshal mocked no-compact mark")
    88  
    89  	markContentReader := strings.NewReader(string(markContent))
    90  	markPath := fmt.Sprintf("%s/phlaredb/%s/%s", userID, meta.ULID.String(), block.NoCompactMarkFilename)
    91  	require.NoError(t, bucket.Upload(context.Background(), markPath, markContentReader))
    92  
    93  	return &mark
    94  }