Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use mockery to generate mocks in data pkg #6244

Merged
merged 1 commit into from
Feb 14, 2025
Merged

use mockery to generate mocks in data pkg #6244

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions flyteadmin/pkg/async/cloudevent/factory_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ func TestGetCloudEventPublisher(t *testing.T) {

db := mocks.NewMockRepository()
mockStore := getMockStore()
url := dataMocks.NewMockRemoteURL()
url := &dataMocks.RemoteURLInterface{}

t.Run("local publisher", func(t *testing.T) {
cfg.Type = common.Local
Expand All @@ -64,7 +64,7 @@ func TestInvalidAwsConfig(t *testing.T) {
}
db := mocks.NewMockRepository()
mockStore := getMockStore()
url := dataMocks.NewMockRemoteURL()
url := &dataMocks.RemoteURLInterface{}

NewCloudEventsPublisher(context.Background(), db, mockStore, url, cfg, remoteCfg, promutils.NewTestScope())
t.Errorf("did not panic")
Expand All @@ -79,7 +79,7 @@ func TestInvalidGcpConfig(t *testing.T) {
}
db := mocks.NewMockRepository()
mockStore := getMockStore()
url := dataMocks.NewMockRemoteURL()
url := &dataMocks.RemoteURLInterface{}

NewCloudEventsPublisher(context.Background(), db, mockStore, url, cfg, remoteCfg, promutils.NewTestScope())
t.Errorf("did not panic")
Expand All @@ -95,7 +95,7 @@ func TestInvalidKafkaConfig(t *testing.T) {
}
db := mocks.NewMockRepository()
mockStore := getMockStore()
url := dataMocks.NewMockRemoteURL()
url := &dataMocks.RemoteURLInterface{}

NewCloudEventsPublisher(context.Background(), db, mockStore, url, cfg, remoteCfg, promutils.NewTestScope())
cfg.KafkaConfig = runtimeInterfaces.KafkaConfig{Version: "2.1.0"}
Expand Down
2 changes: 2 additions & 0 deletions flyteadmin/pkg/data/interfaces/remote.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@ import (
"github.com/flyteorg/flyte/flyteidl/gen/pb-go/flyteidl/admin"
)

//go:generate mockery-v2 --name=RemoteURLInterface --output=../mocks --case=underscore --with-expecter

// Defines an interface for fetching pre-signed URLs.
type RemoteURLInterface interface {
// TODO: Refactor for URI to be of type DataReference. We should package a FromString-like function in flytestdlib
Expand Down
24 changes: 0 additions & 24 deletions flyteadmin/pkg/data/mocks/remote.go
View file
Open in desktop

This file was deleted.

97 changes: 97 additions & 0 deletions flyteadmin/pkg/data/mocks/remote_url_interface.go
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

14 changes: 7 additions & 7 deletions flyteadmin/pkg/manager/impl/execution_manager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -79,7 +79,7 @@ var executionIdentifier = core.WorkflowExecutionIdentifier{
Name: "name",
}
var mockPublisher notificationMocks.Publisher
var mockExecutionRemoteURL = dataMocks.NewMockRemoteURL()
var mockExecutionRemoteURL = &dataMocks.RemoteURLInterface{}
var requestedAt = time.Now()
var testCluster = "C1"
var outputURI = "output uri"
Expand Down Expand Up @@ -3624,8 +3624,8 @@ func TestGetExecutionData(t *testing.T) {
InputsURI: shared.Inputs,
}, nil
}
mockExecutionRemoteURL := dataMocks.NewMockRemoteURL()
mockExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(
mockExecutionRemoteURL := &dataMocks.RemoteURLInterface{}
mockExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(
ctx context.Context, uri string) (*admin.UrlBlob, error) {
if uri == outputURI {
return &admin.UrlBlob{
Expand All @@ -3640,7 +3640,7 @@ func TestGetExecutionData(t *testing.T) {
}

return &admin.UrlBlob{}, errors.New("unexpected input")
}
})
mockStorage := commonMocks.GetMockStorageClient()
fullInputs := &core.LiteralMap{
Literals: map[string]*core.Literal{
Expand Down Expand Up @@ -3842,8 +3842,8 @@ func TestGetExecutionData_LegacyModel(t *testing.T) {
StartedAt: &startedAt,
}, nil
}
mockExecutionRemoteURL := dataMocks.NewMockRemoteURL()
mockExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(
mockExecutionRemoteURL := &dataMocks.RemoteURLInterface{}
mockExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(
ctx context.Context, uri string) (*admin.UrlBlob, error) {
if uri == outputURI {
return &admin.UrlBlob{
Expand All @@ -3858,7 +3858,7 @@ func TestGetExecutionData_LegacyModel(t *testing.T) {
}

return &admin.UrlBlob{}, errors.New("unexpected input")
}
})

repository.ExecutionRepo().(*repositoryMocks.MockExecutionRepo).SetGetCallback(executionGetFunc)
storageClient := getMockStorageForExecTest(context.Background())
Expand Down
9 changes: 5 additions & 4 deletions flyteadmin/pkg/manager/impl/node_execution_manager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"github.com/golang/protobuf/proto"
"github.com/golang/protobuf/ptypes"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/require"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
Expand Down Expand Up @@ -97,7 +98,7 @@ var workflowExecutionIdentifier = core.WorkflowExecutionIdentifier{
Name: "name",
}

var mockNodeExecutionRemoteURL = dataMocks.NewMockRemoteURL()
var mockNodeExecutionRemoteURL = &dataMocks.RemoteURLInterface{}

func addGetExecutionCallback(t *testing.T, repository interfaces.Repository) {
repository.ExecutionRepo().(*repositoryMocks.MockExecutionRepo).SetGetCallback(
Expand Down Expand Up @@ -1377,8 +1378,8 @@ func TestGetNodeExecutionData(t *testing.T) {
}, nil
})

mockNodeExecutionRemoteURL := dataMocks.NewMockRemoteURL()
mockNodeExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
mockNodeExecutionRemoteURL := &dataMocks.RemoteURLInterface{}
mockNodeExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
if uri == "input uri" {
return &admin.UrlBlob{
Url: "inputs",
Expand All @@ -1392,7 +1393,7 @@ func TestGetNodeExecutionData(t *testing.T) {
}

return &admin.UrlBlob{}, errors.New("unexpected input")
}
})
mockStorage := commonMocks.GetMockStorageClient()
fullInputs := &core.LiteralMap{
Literals: map[string]*core.Literal{
Expand Down
9 changes: 5 additions & 4 deletions flyteadmin/pkg/manager/impl/task_execution_manager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (
"github.com/golang/protobuf/proto"
"github.com/golang/protobuf/ptypes"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"
"google.golang.org/grpc/codes"

"github.com/flyteorg/flyte/flyteadmin/pkg/common"
Expand Down Expand Up @@ -63,7 +64,7 @@ var taskEventRequest = &admin.TaskExecutionEventRequest{
},
}

var mockTaskExecutionRemoteURL = dataMocks.NewMockRemoteURL()
var mockTaskExecutionRemoteURL = &dataMocks.RemoteURLInterface{}

var retryAttemptValue = uint32(1)

Expand Down Expand Up @@ -1070,8 +1071,8 @@ func TestGetTaskExecutionData(t *testing.T) {
Closure: closureBytes,
}, nil
})
mockTaskExecutionRemoteURL = dataMocks.NewMockRemoteURL()
mockTaskExecutionRemoteURL.(*dataMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
mockTaskExecutionRemoteURL = &dataMocks.RemoteURLInterface{}
mockTaskExecutionRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
if uri == "input-uri.pb" {
return &admin.UrlBlob{
Url: "inputs",
Expand All @@ -1085,7 +1086,7 @@ func TestGetTaskExecutionData(t *testing.T) {
}

return &admin.UrlBlob{}, errors.New("unexpected input")
}
})
mockStorage := commonMocks.GetMockStorageClient()
fullInputs := &core.LiteralMap{
Literals: map[string]*core.Literal{
Expand Down
19 changes: 10 additions & 9 deletions flyteadmin/pkg/manager/impl/util/data_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (

"github.com/golang/protobuf/proto"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/mock"

"github.com/flyteorg/flyte/flyteadmin/pkg/common"
commonMocks "github.com/flyteorg/flyte/flyteadmin/pkg/common/mocks"
Expand Down Expand Up @@ -118,11 +119,11 @@ func TestGetInputs(t *testing.T) {
Bytes: 1000,
}

mockRemoteURL := urlMocks.NewMockRemoteURL()
mockRemoteURL.(*urlMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
mockRemoteURL := &urlMocks.RemoteURLInterface{}
mockRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
assert.Equal(t, inputsURI, uri)
return expectedURLBlob, nil
}
})
remoteDataConfig := interfaces.RemoteDataConfig{
MaxSizeInBytes: 2000,
}
Expand Down Expand Up @@ -162,11 +163,11 @@ func TestGetOutputs(t *testing.T) {
Bytes: 1000,
}

mockRemoteURL := urlMocks.NewMockRemoteURL()
mockRemoteURL.(*urlMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
mockRemoteURL := &urlMocks.RemoteURLInterface{}
mockRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
assert.Equal(t, testOutputsURI, uri)
return expectedURLBlob, nil
}
})

remoteDataConfig := interfaces.RemoteDataConfig{
MaxSizeInBytes: 2000,
Expand Down Expand Up @@ -205,11 +206,11 @@ func TestGetOutputs(t *testing.T) {
assert.Empty(t, outputURLBlob)
})
t.Run("inline outputs", func(t *testing.T) {
mockRemoteURL := urlMocks.NewMockRemoteURL()
mockRemoteURL.(*urlMocks.MockRemoteURL).GetCallback = func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
mockRemoteURL := &urlMocks.RemoteURLInterface{}
mockRemoteURL.EXPECT().Get(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, uri string) (*admin.UrlBlob, error) {
t.Fatal("Should not fetch a remote URL for outputs stored inline for an execution model")
return &admin.UrlBlob{}, nil
}
})
remoteDataConfig := interfaces.RemoteDataConfig{}
remoteDataConfig.MaxSizeInBytes = 2000

Expand Down
Loading