github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/stores/indexshipper/uploads/table_manager_test.go (about)

     1  package uploads
     2  
     3  import (
     4  	"os"
     5  	"path/filepath"
     6  	"strconv"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/stretchr/testify/require"
    11  
    12  	"github.com/grafana/loki/pkg/storage/chunk/client/local"
    13  	"github.com/grafana/loki/pkg/storage/stores/indexshipper/index"
    14  	"github.com/grafana/loki/pkg/storage/stores/indexshipper/storage"
    15  )
    16  
    17  const objectsStorageDirName = "objects"
    18  
    19  func buildTestStorageClient(t *testing.T, path string) storage.Client {
    20  	objectStoragePath := filepath.Join(path, objectsStorageDirName)
    21  	fsObjectClient, err := local.NewFSObjectClient(local.FSConfig{Directory: objectStoragePath})
    22  	require.NoError(t, err)
    23  
    24  	return storage.NewIndexStorageClient(fsObjectClient, "")
    25  }
    26  
    27  type stopFunc func()
    28  
    29  func buildTestTableManager(t *testing.T, testDir string) (TableManager, stopFunc) {
    30  	storageClient := buildTestStorageClient(t, testDir)
    31  
    32  	cfg := Config{
    33  		UploadInterval: time.Hour,
    34  	}
    35  	tm, err := NewTableManager(cfg, storageClient, nil)
    36  	require.NoError(t, err)
    37  
    38  	return tm, func() {
    39  		tm.Stop()
    40  		require.NoError(t, os.RemoveAll(testDir))
    41  	}
    42  }
    43  
    44  func TestTableManager(t *testing.T) {
    45  	testDir := t.TempDir()
    46  
    47  	testTableManager, stopFunc := buildTestTableManager(t, testDir)
    48  	defer stopFunc()
    49  
    50  	for tableIdx := 0; tableIdx < 2; tableIdx++ {
    51  		tableName := "table-" + strconv.Itoa(tableIdx)
    52  		t.Run(tableName, func(t *testing.T) {
    53  			for userIdx := 0; userIdx < 2; userIdx++ {
    54  				userID := "user-" + strconv.Itoa(userIdx)
    55  				t.Run(userID, func(t *testing.T) {
    56  					userIndexPath := filepath.Join(testDir, tableName, userID)
    57  					require.NoError(t, os.MkdirAll(userIndexPath, 0755))
    58  
    59  					// build some test indexes and add them to the table.
    60  					testIndexes := buildTestIndexes(t, userIndexPath, 5)
    61  					for _, testIndex := range testIndexes {
    62  						require.NoError(t, testTableManager.AddIndex(tableName, userID, testIndex))
    63  					}
    64  
    65  					// see if we can find all the added indexes in the table.
    66  					indexesFound := map[string]*mockIndex{}
    67  					err := testTableManager.ForEach(tableName, userID, func(_ bool, index index.Index) error {
    68  						indexesFound[index.Path()] = index.(*mockIndex)
    69  						return nil
    70  					})
    71  					require.NoError(t, err)
    72  
    73  					require.Equal(t, testIndexes, indexesFound)
    74  				})
    75  			}
    76  		})
    77  	}
    78  }