github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/bucket/tenant_scanner_test.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/storage/tsdb/users_scanner_test.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package bucket
     7  
     8  import (
     9  	"context"
    10  	"errors"
    11  	"path"
    12  	"testing"
    13  
    14  	"github.com/go-kit/log"
    15  	"github.com/stretchr/testify/assert"
    16  	"github.com/stretchr/testify/require"
    17  
    18  	"github.com/grafana/pyroscope/pkg/test/mocks/mockobjstore"
    19  )
    20  
    21  func TestUsersScanner_ScanUsers_ShouldReturnedOwnedUsersOnly(t *testing.T) {
    22  	bucketClient := mockobjstore.NewMockBucketWithHelper(t)
    23  	bucketClient.MockIter("", []string{"user-1", "user-2", "user-3", "user-4"}, nil)
    24  	bucketClient.MockExists(path.Join("user-1", "phlaredb/", TenantDeletionMarkPath), false, nil)
    25  	bucketClient.MockExists(path.Join("user-3", "phlaredb/", TenantDeletionMarkPath), true, nil)
    26  
    27  	isOwned := func(userID string) (bool, error) {
    28  		return userID == "user-1" || userID == "user-3", nil
    29  	}
    30  
    31  	s := NewTenantsScanner(bucketClient, isOwned, log.NewNopLogger())
    32  	actual, deleted, err := s.ScanTenants(context.Background())
    33  	require.NoError(t, err)
    34  	assert.Equal(t, []string{"user-1"}, actual)
    35  	assert.Equal(t, []string{"user-3"}, deleted)
    36  }
    37  
    38  func TestUsersScanner_ScanUsers_ShouldReturnUsersForWhichOwnerCheckOrTenantDeletionCheckFailed(t *testing.T) {
    39  	expected := []string{"user-1", "user-2"}
    40  
    41  	bucketClient := mockobjstore.NewMockBucketWithHelper(t)
    42  	bucketClient.MockIter("", expected, nil)
    43  	bucketClient.MockExists(path.Join("user-1", "phlaredb/", TenantDeletionMarkPath), false, nil)
    44  	bucketClient.MockExists(path.Join("user-2", "phlaredb/", TenantDeletionMarkPath), false, errors.New("fail"))
    45  
    46  	isOwned := func(userID string) (bool, error) {
    47  		return false, errors.New("failed to check if user is owned")
    48  	}
    49  
    50  	s := NewTenantsScanner(bucketClient, isOwned, log.NewNopLogger())
    51  	actual, deleted, err := s.ScanTenants(context.Background())
    52  	require.NoError(t, err)
    53  	assert.Equal(t, expected, actual)
    54  	assert.Empty(t, deleted)
    55  }
    56  
    57  func TestUsersScanner_ScanUsers_ShouldNotReturnPrefixedUsedByPyroscopeInternals(t *testing.T) {
    58  	bucketClient := mockobjstore.NewMockBucketWithHelper(t)
    59  	bucketClient.MockIter("", []string{"user-1", "user-2", PyroscopeInternalsPrefix}, nil)
    60  	bucketClient.MockExists(path.Join("user-1", "phlaredb/", TenantDeletionMarkPath), false, nil)
    61  	bucketClient.MockExists(path.Join("user-2", "phlaredb/", TenantDeletionMarkPath), false, nil)
    62  
    63  	s := NewTenantsScanner(bucketClient, AllTenants, log.NewNopLogger())
    64  	actual, _, err := s.ScanTenants(context.Background())
    65  	require.NoError(t, err)
    66  	assert.Equal(t, []string{"user-1", "user-2"}, actual)
    67  }