drone/registry/app/pkg/python/local_helper_test.go

// Copyright 2023 Harness, Inc.
//
// Licensed 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 python

import (
	"context"
	"errors"
	"io"
	"mime/multipart"
	"testing"

	"github.com/harness/gitness/registry/app/api/openapi/contracts/artifact"
	"github.com/harness/gitness/registry/app/metadata"
	"github.com/harness/gitness/registry/app/pkg"
	"github.com/harness/gitness/registry/app/pkg/commons"
	"github.com/harness/gitness/registry/app/pkg/types/python"
	"github.com/harness/gitness/registry/app/storage"
	"github.com/harness/gitness/registry/types"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/mock"
)

type MockLocalRegistry struct {
	mock.Mock
}

func (m *MockLocalRegistry) GetArtifactType() artifact.RegistryType {
	args := m.Called()
	return args.Get(0).(artifact.RegistryType) //nolint:errcheck
}

func (m *MockLocalRegistry) GetPackageTypes() []artifact.PackageType {
	args := m.Called()
	return args.Get(0).([]artifact.PackageType) //nolint:errcheck
}

func (m *MockLocalRegistry) DownloadPackageFile(ctx context.Context, info python.ArtifactInfo) (
	*commons.ResponseHeaders, *storage.FileReader, io.ReadCloser, string, error,
) {
	args := m.Called(ctx, info)
	//nolint:errcheck
	return args.Get(0).(*commons.ResponseHeaders), args.Get(1).(*storage.FileReader),
		args.Get(2).(io.ReadCloser), args.String(3), args.Error(4)
}

func (m *MockLocalRegistry) GetPackageMetadata(ctx context.Context, info python.ArtifactInfo) (
	python.PackageMetadata,
	error,
) {
	args := m.Called(ctx, info)
	return args.Get(0).(python.PackageMetadata), args.Error(1) //nolint:errcheck
}

func (m *MockLocalRegistry) UploadPackageFile(
	ctx context.Context,
	info python.ArtifactInfo,
	file multipart.File,
	filename string,
) (*commons.ResponseHeaders, string, error) {
	args := m.Called(ctx, info, file, filename)
	return args.Get(0).(*commons.ResponseHeaders), args.String(1), args.Error(2) //nolint:errcheck
}

func (m *MockLocalRegistry) UploadPackageFileReader(
	ctx context.Context,
	info python.ArtifactInfo,
	file io.ReadCloser,
	filename string,
) (*commons.ResponseHeaders, string, error) {
	args := m.Called(ctx, info, file, filename)
	return args.Get(0).(*commons.ResponseHeaders), args.String(1), args.Error(2) //nolint:errcheck
}

type MockLocalBase struct {
	mock.Mock
}

func (m *MockLocalBase) MoveTempFile(
	_ context.Context,
	_ pkg.ArtifactInfo,
	_, _, _ string,
	_ metadata.Metadata,
	_ types.FileInfo,
) (*commons.ResponseHeaders, string, error) {
	// TODO implement me
	panic("implement me")
}

func (m *MockLocalBase) CheckIfVersionExists(_ context.Context, _ pkg.PackageArtifactInfo) (bool, error) {
	// TODO implement me
	panic("implement me")
}

func (m *MockLocalBase) DeletePackage(_ context.Context, _ pkg.PackageArtifactInfo) error {
	// TODO implement me
	panic("implement me")
}

func (m *MockLocalBase) DeleteVersion(_ context.Context, _ pkg.PackageArtifactInfo) error {
	// TODO implement me
	panic("implement me")
}

func (m *MockLocalBase) Exists(ctx context.Context, info pkg.ArtifactInfo, version, filename string) bool {
	args := m.Called(ctx, info, version, filename)
	return args.Bool(0)
}

func (m *MockLocalBase) ExistsByFilePath(ctx context.Context, registryID int64, filePath string) (bool, error) {
	args := m.Called(ctx, registryID, filePath)
	return args.Bool(0), args.Error(1)
}

func (m *MockLocalBase) Download(ctx context.Context, info pkg.ArtifactInfo, version, filename string) (
	*commons.ResponseHeaders, *storage.FileReader, string, error,
) {
	args := m.Called(ctx, info, version, filename)
	//nolint:errcheck
	return args.Get(0).(*commons.ResponseHeaders), args.Get(1).(*storage.FileReader),
		args.String(2), args.Error(3)
}

func (m *MockLocalBase) Upload(
	ctx context.Context, info pkg.ArtifactInfo, filename, version, path string,
	reader io.ReadCloser, metadata metadata.Metadata,
) (*commons.ResponseHeaders, string, error) {
	args := m.Called(ctx, info, filename, version, path, reader, metadata)
	return args.Get(0).(*commons.ResponseHeaders), args.String(1), args.Error(2) //nolint:errcheck
}

func (m *MockLocalBase) Move(
	ctx context.Context,
	info pkg.ArtifactInfo,
	filename, version, path string,
	metadata metadata.Metadata,
	fileInfo types.FileInfo,
) (*commons.ResponseHeaders, string, error) {
	args := m.Called(ctx, info, filename, version, path, metadata, fileInfo)
	return args.Get(0).(*commons.ResponseHeaders), args.String(1), args.Error(2) //nolint:errcheck
}

func (m *MockLocalBase) UploadFile(
	ctx context.Context, info pkg.ArtifactInfo, filename, version, path string,
	file multipart.File,
	metadata metadata.Metadata,
) (*commons.ResponseHeaders, string, error) {
	args := m.Called(ctx, info, filename, version, path, file, metadata)
	return args.Get(0).(*commons.ResponseHeaders), args.String(1), args.Error(2) //nolint:errcheck
}

type MockReadCloser struct {
	mock.Mock
}

func (m *MockReadCloser) Read(p []byte) (n int, err error) {
	args := m.Called(p)
	return args.Int(0), args.Error(1)
}

func (m *MockReadCloser) Close() error {
	args := m.Called()
	return args.Error(0)
}

func TestNewLocalRegistryHelper(t *testing.T) {
	mockLocalRegistry := new(MockLocalRegistry)
	mockLocalBase := new(MockLocalBase)

	helper := NewLocalRegistryHelper(mockLocalRegistry, mockLocalBase)

	assert.NotNil(t, helper, "Helper should not be nil")
}

// Test for FileExists method.
func TestLocalRegistryHelper_FileExists(t *testing.T) {
	mockLocalRegistry := new(MockLocalRegistry)
	mockLocalBase := new(MockLocalBase)

	helper := NewLocalRegistryHelper(mockLocalRegistry, mockLocalBase)

	ctx := context.Background()
	artifactInfo := python.ArtifactInfo{
		ArtifactInfo: pkg.ArtifactInfo{
			BaseInfo: &pkg.BaseInfo{
				RootParentID:   1,
				RootIdentifier: "root",
				//RegParentID:    2,
			},
			RegIdentifier: "registry",
			Image:         "package",
		},
		Version:  "1.0.0",
		Filename: "package-1.0.0.whl",
	}

	mockLocalBase.On("Exists", ctx, artifactInfo.ArtifactInfo, artifactInfo.Version, artifactInfo.Filename).Return(true)

	exists := helper.FileExists(ctx, artifactInfo)

	assert.True(t, exists, "File should exist")
	mockLocalBase.AssertExpectations(t)
}

// Test for DownloadFile method.
func TestLocalRegistryHelper_DownloadFile(t *testing.T) {
	mockLocalRegistry := new(MockLocalRegistry)
	mockLocalBase := new(MockLocalBase)

	helper := NewLocalRegistryHelper(mockLocalRegistry, mockLocalBase)

	ctx := context.Background()

	artifactInfo := python.ArtifactInfo{
		ArtifactInfo: pkg.ArtifactInfo{
			BaseInfo: &pkg.BaseInfo{
				RootParentID:   1,
				RootIdentifier: "root",
				//RegParentID:    2,
			},
			RegIdentifier: "registry",
			Image:         "package",
		},
		Version:  "1.0.0",
		Filename: "package-1.0.0.whl",
	}

	expectedHeaders := &commons.ResponseHeaders{}
	expectedFileReader := &storage.FileReader{}
	expectedRedirectURL := "http://example.com/download"
	var expectedError error

	mockLocalBase.On("Download", ctx, artifactInfo.ArtifactInfo, artifactInfo.Version, artifactInfo.Filename).
		Return(expectedHeaders, expectedFileReader, expectedRedirectURL, expectedError)

	headers, fileReader, redirectURL, err := helper.DownloadFile(ctx, artifactInfo)

	assert.Equal(t, expectedHeaders, headers, "Headers should match")
	assert.Equal(t, expectedFileReader, fileReader, "FileReader should match")
	assert.Equal(t, expectedRedirectURL, redirectURL, "RedirectURL should match")
	assert.Nil(t, err, "Error should be nil")
	mockLocalBase.AssertExpectations(t)
}

// Test for DownloadFile method with error.
func TestLocalRegistryHelper_DownloadFile_Error(t *testing.T) {
	mockLocalRegistry := new(MockLocalRegistry)
	mockLocalBase := new(MockLocalBase)
	helper := NewLocalRegistryHelper(mockLocalRegistry, mockLocalBase)

	ctx := context.Background()
	artifactInfo := python.ArtifactInfo{
		ArtifactInfo: pkg.ArtifactInfo{
			BaseInfo: &pkg.BaseInfo{
				RootParentID:   1,
				RootIdentifier: "root",
				//RegParentID:    2,
			},
			RegIdentifier: "registry",
			Image:         "package",
		},
		Version:  "1.0.0",
		Filename: "package-1.0.0.whl",
	}

	expectedError := errors.New("download error")

	mockLocalBase.On("Download", ctx, artifactInfo.ArtifactInfo, artifactInfo.Version, artifactInfo.Filename).
		Return((*commons.ResponseHeaders)(nil), (*storage.FileReader)(nil), "", expectedError)

	headers, fileReader, redirectURL, err := helper.DownloadFile(ctx, artifactInfo)

	assert.Nil(t, headers, "Headers should be nil")
	assert.Nil(t, fileReader, "FileReader should be nil")
	assert.Equal(t, "", redirectURL, "RedirectURL should be empty")
	assert.Equal(t, expectedError, err, "Error should match expected error")
	mockLocalBase.AssertExpectations(t)
}

// Test for UploadPackageFile method.
func TestLocalRegistryHelper_UploadPackageFile(t *testing.T) {
	mockLocalRegistry := new(MockLocalRegistry)
	mockLocalBase := new(MockLocalBase)
	helper := NewLocalRegistryHelper(mockLocalRegistry, mockLocalBase)

	ctx := context.Background()
	artifactInfo := python.ArtifactInfo{
		ArtifactInfo: pkg.ArtifactInfo{
			BaseInfo: &pkg.BaseInfo{
				RootParentID:   1,
				RootIdentifier: "root",
				//RegParentID:    2,
			},
			RegIdentifier: "registry",
			Image:         "package",
		},
		Version:  "1.0.0",
		Filename: "package-1.0.0.whl",
	}

	mockReader := new(MockReadCloser)
	mockReader.On("Close").Return(nil)

	expectedHeaders := &commons.ResponseHeaders{}
	expectedSHA256 := "abc123"
	var expectedError error

	mockLocalRegistry.On("UploadPackageFileReader", ctx, artifactInfo, mock.AnythingOfType("*python.MockReadCloser"),
		"package-1.0.0.whl").
		Return(expectedHeaders, expectedSHA256, expectedError)

	headers, sha256, err := helper.UploadPackageFile(ctx, artifactInfo, mockReader, "package-1.0.0.whl")

	assert.Equal(t, expectedHeaders, headers, "Headers should match")
	assert.Equal(t, expectedSHA256, sha256, "SHA256 should match")
	assert.Nil(t, err, "Error should be nil")
	mockLocalRegistry.AssertExpectations(t)
}

// Test for UploadPackageFile method with error.
func TestLocalRegistryHelper_UploadPackageFile_Error(t *testing.T) {
	mockLocalRegistry := new(MockLocalRegistry)
	mockLocalBase := new(MockLocalBase)

	helper := NewLocalRegistryHelper(mockLocalRegistry, mockLocalBase)

	ctx := context.Background()
	artifactInfo := python.ArtifactInfo{
		ArtifactInfo: pkg.ArtifactInfo{
			BaseInfo: &pkg.BaseInfo{
				RootParentID:   1,
				RootIdentifier: "root",
				//RegParentID:    2,
			},
			RegIdentifier: "registry",
			Image:         "package",
		},
		Version:  "1.0.0",
		Filename: "package-1.0.0.whl",
	}

	mockReader := new(MockReadCloser)
	mockReader.On("Close").Return(nil)

	expectedError := errors.New("upload error")

	mockLocalRegistry.On("UploadPackageFileReader", ctx, artifactInfo, mock.AnythingOfType("*python.MockReadCloser"),
		"package-1.0.0.whl").
		Return((*commons.ResponseHeaders)(nil), "", expectedError)

	headers, sha256, err := helper.UploadPackageFile(ctx, artifactInfo, mockReader, "package-1.0.0.whl")

	assert.Nil(t, headers, "Headers should be nil")
	assert.Equal(t, "", sha256, "SHA256 should be empty")
	assert.Equal(t, expectedError, err, "Error should match expected error")
	mockLocalRegistry.AssertExpectations(t)
}