github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/externalresource/internal/bucket/s3_test_utils.go (about)

     1  // Copyright 2022 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package bucket
    15  
    16  import (
    17  	"context"
    18  	"fmt"
    19  	"os"
    20  	"path/filepath"
    21  	"testing"
    22  
    23  	brStorage "github.com/pingcap/tidb/br/pkg/storage"
    24  	"github.com/pingcap/tiflow/engine/model"
    25  	"github.com/pingcap/tiflow/engine/pkg/externalresource/internal"
    26  	resModel "github.com/pingcap/tiflow/engine/pkg/externalresource/model"
    27  	"github.com/pingcap/tiflow/pkg/errors"
    28  	"github.com/stretchr/testify/require"
    29  )
    30  
    31  const (
    32  	envS3Endpoint        = "ENGINE_S3_ENDPOINT"
    33  	envS3AccessKeyID     = "ENGINE_S3_ACCESS_KEY"
    34  	envS3SecretAccessKey = "ENGINE_S3_SECRET_KEY"
    35  
    36  	// UtBucketName is the bucket name for UT
    37  	UtBucketName = "engine-ut"
    38  	// MockExecutorID is the executor ID for mock executor
    39  	MockExecutorID = "executor-1"
    40  )
    41  
    42  // GetS3OptionsForUT returns the s3 options for unit test.
    43  func GetS3OptionsForUT() (*brStorage.S3BackendOptions, error) {
    44  	endpoint := os.Getenv(envS3Endpoint)
    45  	if len(endpoint) == 0 {
    46  		return nil, errors.Errorf("empty endpoint in env %s", envS3Endpoint)
    47  	}
    48  
    49  	accessKeyID := os.Getenv(envS3AccessKeyID)
    50  	if len(accessKeyID) == 0 {
    51  		return nil, errors.Errorf("empty access key ID in env %s", envS3AccessKeyID)
    52  	}
    53  
    54  	secretAccessKey := os.Getenv(envS3SecretAccessKey)
    55  	if len(secretAccessKey) == 0 {
    56  		return nil, errors.Errorf("empty secret access key in env %s", envS3SecretAccessKey)
    57  	}
    58  
    59  	return &brStorage.S3BackendOptions{
    60  		Endpoint:        endpoint,
    61  		AccessKey:       accessKeyID,
    62  		SecretAccessKey: secretAccessKey,
    63  		Provider:        "minio",
    64  		ForcePathStyle:  true,
    65  	}, nil
    66  }
    67  
    68  type mockBucketCreator struct {
    69  	baseDir string
    70  	bucket  string
    71  }
    72  
    73  func newMockBucketCreator(tempDir string, bucket string) *mockBucketCreator {
    74  	return &mockBucketCreator{
    75  		baseDir: tempDir,
    76  		bucket:  bucket,
    77  	}
    78  }
    79  
    80  func (f *mockBucketCreator) newBucketForScope(
    81  	ctx context.Context, scope internal.ResourceScope,
    82  ) (brStorage.ExternalStorage, error) {
    83  	uri := fmt.Sprintf("%s/%s", f.baseURI(), scope.BuildResPath())
    84  	return f.newBucketFromURI(ctx, uri)
    85  }
    86  
    87  func (f *mockBucketCreator) newBucketFromURI(
    88  	ctx context.Context,
    89  	uri string,
    90  ) (brStorage.ExternalStorage, error) {
    91  	return brStorage.NewLocalStorage(uri)
    92  }
    93  
    94  func (f *mockBucketCreator) baseURI() string {
    95  	return fmt.Sprintf("%s/%s", f.baseDir, f.bucket)
    96  }
    97  
    98  func (f *mockBucketCreator) resourceType() resModel.ResourceType {
    99  	return resModel.ResourceTypeS3
   100  }
   101  
   102  func (f *mockBucketCreator) assertFileExists(t *testing.T, uri string) {
   103  	require.FileExists(t, filepath.Join(f.baseDir, uri))
   104  }
   105  
   106  func (f *mockBucketCreator) assertFileNotExist(t *testing.T, uri string) {
   107  	require.NoFileExists(t, filepath.Join(f.baseDir, uri))
   108  }
   109  
   110  // NewFileManagerForUT returns a file manager for UT.
   111  func NewFileManagerForUT(tempDir string, executorID resModel.ExecutorID) (*FileManager, *mockBucketCreator) {
   112  	creator := newMockBucketCreator(tempDir, UtBucketName)
   113  	return NewFileManager(
   114  		executorID,
   115  		creator,
   116  	), creator
   117  }
   118  
   119  // NewFileManagerForUTFromSharedStorageFactory returns a file manager for UT.
   120  func NewFileManagerForUTFromSharedStorageFactory(
   121  	executorID model.ExecutorID, creator *mockBucketCreator,
   122  ) *FileManager {
   123  	return NewFileManager(executorID, creator)
   124  }