github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/engine/pkg/externalresource/broker/storage_handle_test.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 broker
    15  
    16  import (
    17  	"context"
    18  	"testing"
    19  
    20  	pb "github.com/pingcap/tiflow/engine/enginepb"
    21  	"github.com/pingcap/tiflow/engine/pkg/externalresource/internal"
    22  	"github.com/pingcap/tiflow/engine/pkg/externalresource/internal/local"
    23  	"github.com/pingcap/tiflow/engine/pkg/externalresource/manager"
    24  	resModel "github.com/pingcap/tiflow/engine/pkg/externalresource/model"
    25  	"github.com/pingcap/tiflow/engine/pkg/tenant"
    26  	"github.com/stretchr/testify/mock"
    27  	"github.com/stretchr/testify/require"
    28  )
    29  
    30  func newResourceIdentForTesting(executor, workerID, resourceName string) internal.ResourceIdent {
    31  	return internal.ResourceIdent{
    32  		Name: resModel.EncodeResourceName(resourceName),
    33  		ResourceScope: internal.ResourceScope{
    34  			ProjectInfo: tenant.NewProjectInfo("fakeTenant", "fakeProject"),
    35  			Executor:    resModel.ExecutorID(executor),
    36  			WorkerID:    workerID,
    37  		},
    38  	}
    39  }
    40  
    41  func TestStorageHandlePersistAndDiscard(t *testing.T) {
    42  	t.Parallel()
    43  	dir := t.TempDir()
    44  
    45  	executor := resModel.ExecutorID("executor-1")
    46  	ident := newResourceIdentForTesting(string(executor), "worker-1", "test-resource")
    47  	fm := local.NewLocalFileManager(executor, resModel.LocalFileConfig{BaseDir: dir})
    48  	cli := manager.NewMockClient()
    49  
    50  	ctx := context.Background()
    51  	desc, err := fm.CreateResource(ctx, ident)
    52  	require.NoError(t, err)
    53  
    54  	handle, err := newResourceHandle(
    55  		"job-1",
    56  		executor,
    57  		fm, desc, false, cli)
    58  	require.NoError(t, err)
    59  
    60  	cli.On("CreateResource", mock.Anything, &pb.CreateResourceRequest{
    61  		ProjectInfo: &pb.ProjectInfo{
    62  			TenantId:  desc.ResourceIdent().TenantID(),
    63  			ProjectId: desc.ResourceIdent().ProjectID(),
    64  		},
    65  		ResourceId:      "/local/test-resource",
    66  		CreatorExecutor: string(executor),
    67  		JobId:           "job-1",
    68  		CreatorWorkerId: "worker-1",
    69  	}).Return(nil).Once()
    70  	err = handle.Persist(context.Background())
    71  	require.NoError(t, err)
    72  	cli.AssertExpectations(t)
    73  	cli.ExpectedCalls = nil
    74  
    75  	err = handle.Persist(context.Background())
    76  	require.NoError(t, err)
    77  
    78  	desc, err = fm.GetPersistedResource(ctx, ident)
    79  	require.NoError(t, err)
    80  	require.NotNil(t, desc)
    81  
    82  	cli.On("RemoveResource", mock.Anything, &pb.RemoveResourceRequest{
    83  		ResourceKey: &pb.ResourceKey{
    84  			JobId:      "job-1",
    85  			ResourceId: "/local/test-resource",
    86  		},
    87  	}).Return(nil).Once()
    88  	err = handle.Discard(context.Background())
    89  	require.NoError(t, err)
    90  	cli.AssertExpectations(t)
    91  	cli.ExpectedCalls = nil
    92  
    93  	_, err = fm.GetPersistedResource(ctx, ident)
    94  	require.Error(t, err)
    95  	require.Regexp(t, ".*ErrResourceDoesNotExist.*", err)
    96  
    97  	// Discarding twice should fail.
    98  	err = handle.Discard(context.Background())
    99  	require.Error(t, err)
   100  	require.Regexp(t, ".*ErrInvalidResourceHandle.*", err)
   101  
   102  	// Persisting after discarding should fail.
   103  	err = handle.Persist(context.Background())
   104  	require.Error(t, err)
   105  	require.Regexp(t, ".*ErrInvalidResourceHandle.*", err)
   106  }
   107  
   108  func TestStorageHandleDiscardTemporaryResource(t *testing.T) {
   109  	t.Parallel()
   110  	dir := t.TempDir()
   111  	fm := local.NewLocalFileManager("", resModel.LocalFileConfig{BaseDir: dir})
   112  	cli := manager.NewMockClient()
   113  
   114  	ctx := context.Background()
   115  	desc, err := fm.CreateResource(ctx, newResourceIdentForTesting("", "worker-1", "test-resource"))
   116  	require.NoError(t, err)
   117  
   118  	handle, err := newResourceHandle(
   119  		"job-1",
   120  		"executor-1",
   121  		fm, desc, false, cli)
   122  	require.NoError(t, err)
   123  
   124  	err = handle.Discard(context.Background())
   125  	require.NoError(t, err)
   126  	cli.AssertNotCalled(t, "RemoveResource")
   127  	cli.ExpectedCalls = nil
   128  
   129  	_, err = fm.GetPersistedResource(ctx, newResourceIdentForTesting("", "worker-1", "test-resource"))
   130  	require.Error(t, err)
   131  	require.Regexp(t, ".*ErrResourceDoesNotExist.*", err)
   132  
   133  	err = handle.Discard(context.Background())
   134  	require.Error(t, err)
   135  	require.Regexp(t, ".*ErrInvalidResourceHandle.*", err)
   136  }