milvus/internal/util/fileresource/manager_test.go

/*
 * # Licensed to the LF AI & Data foundation under one
 * # or more contributor license agreements. See the NOTICE file
 * # distributed with this work for additional information
 * # regarding copyright ownership. The ASF licenses this file
 * # to you under the Apache License, Version 2.0 (the
 * # "License"); you may not use this file except in compliance
 * # with the License. You may obtain a copy of the License at
 * #
 * #     http://www.apache.org/licenses/LICENSE-2.0
 * #
 * # Unless required by applicable law or agreed to in writing, software
 * # distributed under the License is distributed on an "AS IS" BASIS,
 * # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * # See the License for the specific language governing permissions and
 * # limitations under the License.
 */

package fileresource

import (
	"context"
	"fmt"
	"io"
	"os"
	"path"
	"strings"
	"sync"
	"testing"

	"github.com/stretchr/testify/suite"

	"github.com/milvus-io/milvus/internal/mocks"
	"github.com/milvus-io/milvus/pkg/v2/proto/internalpb"
	"github.com/milvus-io/milvus/pkg/v2/util/conc"
	"github.com/milvus-io/milvus/pkg/v2/util/typeutil"
)

// MockReader implements storage.FileReader using composition
type MockReader struct {
	io.Reader
	io.Closer
	io.ReaderAt
	io.Seeker
	size int64
}

func (mr *MockReader) Size() (int64, error) {
	return mr.size, nil
}

func newMockReader(s string) *MockReader {
	ioReader := strings.NewReader(s)
	return &MockReader{Reader: ioReader, Closer: io.NopCloser(ioReader), size: int64(len(s))}
}

// BaseManagerSuite tests BaseManager
type BaseManagerSuite struct {
	suite.Suite
	manager *BaseManager
}

func (suite *BaseManagerSuite) SetupTest() {
	suite.manager = &BaseManager{localPath: "/tmp/test"}
}

func (suite *BaseManagerSuite) TestSync() {
	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test.file", Path: "/test/test.file"},
	}

	err := suite.manager.Sync(resources)
	suite.NoError(err)
}

func (suite *BaseManagerSuite) TestDownload() {
	mockStorage := mocks.NewChunkManager(suite.T())
	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test.file", Path: "/test/test.file"},
	}

	err := suite.manager.Download(mockStorage, resources...)
	suite.NoError(err)
}

func (suite *BaseManagerSuite) TestRelease() {
	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test.file", Path: "/test/test.file"},
	}

	suite.manager.Release(resources...)
	// Should not panic
}

func (suite *BaseManagerSuite) TestMode() {
	mode := suite.manager.Mode()
	suite.Equal(CloseMode, mode)
}

func TestBaseManagerSuite(t *testing.T) {
	suite.Run(t, new(BaseManagerSuite))
}

// SyncManagerSuite tests SyncManager
type SyncManagerSuite struct {
	suite.Suite
	manager     *SyncManager
	mockStorage *mocks.ChunkManager
	tempDir     string
}

func (suite *SyncManagerSuite) SetupTest() {
	// Create temporary directory for tests
	var err error
	suite.tempDir, err = os.MkdirTemp(os.TempDir(), "fileresource_test_sync_*")
	suite.NoError(err)

	suite.mockStorage = mocks.NewChunkManager(suite.T())
	suite.manager = &SyncManager{
		BaseManager: BaseManager{localPath: suite.tempDir},
		downloader:  suite.mockStorage,
		resourceSet: make(map[int64]struct{}),
	}
}

func (suite *SyncManagerSuite) TearDownTest() {
	if suite.tempDir != "" {
		os.RemoveAll(suite.tempDir)
	}
}

func (suite *SyncManagerSuite) TestSync_Success() {
	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test1.file", Path: "/storage/test1.file"},
		{Id: 2, Name: "test2.file", Path: "/storage/test2.file"},
	}

	// Mock the Reader calls
	suite.mockStorage.EXPECT().Reader(context.Background(), "/storage/test1.file").Return(newMockReader("test content 1"), nil)
	suite.mockStorage.EXPECT().Reader(context.Background(), "/storage/test2.file").Return(newMockReader("test content 2"), nil)

	err := suite.manager.Sync(resources)
	suite.NoError(err)

	// Verify files were created
	file1Path := path.Join(suite.tempDir, "1", "test1.file")
	file2Path := path.Join(suite.tempDir, "2", "test2.file")

	suite.FileExists(file1Path)
	suite.FileExists(file2Path)

	// Verify content
	content1, err := os.ReadFile(file1Path)
	suite.NoError(err)
	suite.Equal("test content 1", string(content1))

	content2, err := os.ReadFile(file2Path)
	suite.NoError(err)
	suite.Equal("test content 2", string(content2))
}

func (suite *SyncManagerSuite) TestSync_ReaderError() {
	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test.file", Path: "/storage/nonexistent.file"},
	}

	// Mock reader to return error
	suite.mockStorage.EXPECT().Reader(context.Background(), "/storage/nonexistent.file").Return(nil, io.ErrUnexpectedEOF)

	err := suite.manager.Sync(resources)
	suite.Error(err)
	suite.ErrorIs(err, io.ErrUnexpectedEOF)
}

func (suite *SyncManagerSuite) TestMode() {
	mode := suite.manager.Mode()
	suite.Equal(SyncMode, mode)
}

func TestSyncManagerSuite(t *testing.T) {
	suite.Run(t, new(SyncManagerSuite))
}

// RefManagerSuite tests RefManager
type RefManagerSuite struct {
	suite.Suite
	manager     *RefManager
	mockStorage *mocks.ChunkManager
	tempDir     string
}

func (suite *RefManagerSuite) SetupTest() {
	// Create temporary directory for tests
	var err error
	suite.tempDir, err = os.MkdirTemp(os.TempDir(), "fileresource_test_ref_*")
	suite.NoError(err)

	suite.mockStorage = mocks.NewChunkManager(suite.T())
	suite.manager = &RefManager{
		BaseManager: BaseManager{localPath: suite.tempDir},
		ref:         map[string]int{},
		finished:    typeutil.NewConcurrentMap[string, bool](),
		sf:          &conc.Singleflight[interface{}]{},
	}
}

func (suite *RefManagerSuite) TearDownTest() {
	if suite.tempDir != "" {
		os.RemoveAll(suite.tempDir)
	}
}

func (suite *RefManagerSuite) TestNormal() {
	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test", Path: "/storage/test.file"},
	}

	// Set up mock
	suite.mockStorage.EXPECT().RootPath().Return("/test/storage")
	suite.mockStorage.EXPECT().Reader(context.Background(), "/storage/test.file").Return(newMockReader("test content"), nil)

	err := suite.manager.Download(suite.mockStorage, resources...)
	suite.Require().NoError(err)

	// Verify ref count
	key := "/test/storage/1"
	suite.Equal(1, suite.manager.ref[key])

	// Verify storage name is set
	suite.Equal("/test/storage", resources[0].StorageName)

	// Verify file was downloaded
	// {local_path}/{storage_name}/{resource_id}/{file_name}
	filePath := path.Join(suite.tempDir, "/test/storage", fmt.Sprint(1), path.Base(resources[0].GetPath()))
	suite.FileExists(filePath)

	content, err := os.ReadFile(filePath)
	suite.NoError(err)
	suite.Equal("test content", string(content))

	// release and clean all file
	suite.manager.Release(resources...)
	suite.manager.CleanResource()
	suite.NoFileExists(filePath)
}

func (suite *RefManagerSuite) TestMode() {
	mode := suite.manager.Mode()
	suite.Equal(RefMode, mode)
}

func TestRefManagerSuite(t *testing.T) {
	suite.Run(t, new(RefManagerSuite))
}

// ManagerFactorySuite tests NewManager factory function
type ManagerFactorySuite struct {
	suite.Suite
	mockStorage *mocks.ChunkManager
}

func (suite *ManagerFactorySuite) SetupTest() {
	suite.mockStorage = mocks.NewChunkManager(suite.T())
}

func (suite *ManagerFactorySuite) TestNewManager_BaseManager() {
	manager := NewManager(suite.mockStorage, CloseMode)
	suite.IsType(&BaseManager{}, manager)
	suite.Equal(CloseMode, manager.Mode())
}

func (suite *ManagerFactorySuite) TestNewManager_SyncManager() {
	manager := NewManager(suite.mockStorage, SyncMode)
	suite.IsType(&SyncManager{}, manager)
	suite.Equal(SyncMode, manager.Mode())
}

func (suite *ManagerFactorySuite) TestNewManager_InvalidMode() {
	suite.Panics(func() {
		NewManager(suite.mockStorage, Mode(999))
	})
}

func TestManagerFactorySuite(t *testing.T) {
	suite.Run(t, new(ManagerFactorySuite))
}

// GlobalFunctionsSuite tests global functions
type GlobalFunctionsSuite struct {
	suite.Suite
	mockStorage *mocks.ChunkManager
}

func (suite *GlobalFunctionsSuite) SetupTest() {
	suite.mockStorage = mocks.NewChunkManager(suite.T())
	// Reset global state
	GlobalFileManager = nil
	once = sync.Once{}
}

func (suite *GlobalFunctionsSuite) TestInitManager() {
	InitManager(suite.mockStorage, SyncMode)

	suite.NotNil(GlobalFileManager)
	suite.Equal(SyncMode, GlobalFileManager.Mode())

	// Test that calling InitManager again doesn't change the manager
	oldManager := GlobalFileManager
	InitManager(suite.mockStorage, RefMode)
	suite.Equal(oldManager, GlobalFileManager)
	suite.Equal(SyncMode, GlobalFileManager.Mode()) // Should still be SyncMode
}

func (suite *GlobalFunctionsSuite) TestSync_NotInitialized() {
	GlobalFileManager = nil

	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test.file", Path: "/test/test.file"},
	}

	err := Sync(resources)
	suite.NoError(err) // Should not error when not initialized
}

func (suite *GlobalFunctionsSuite) TestSync_Initialized() {
	InitManager(suite.mockStorage, CloseMode)

	resources := []*internalpb.FileResourceInfo{
		{Id: 1, Name: "test.file", Path: "/test/test.file"},
	}

	err := Sync(resources)
	suite.NoError(err)
}

func TestGlobalFunctionsSuite(t *testing.T) {
	suite.Run(t, new(GlobalFunctionsSuite))
}