Skip to content

Commit

Permalink
Fix new uses of mocks
Browse files Browse the repository at this point in the history 
Signed-off-by: Eduardo Apolinario <[email protected]>
  • Loading branch information
@eapolinario
eapolinario committed Feb 27, 2025
1 parent c2b4423 commit 098c2b1
Show file tree
Hide file tree
Showing 3 changed files with 14 additions and 14 deletions.
14 changes: 7 additions & 7 deletions flytectl/cmd/register/files_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ func TestRegisterFromFiles(t *testing.T) {
s.MockAdminClient.EXPECT().CreateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().UpdateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
mockDataProxy := s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient)
mockDataProxy.OnCreateUploadLocationMatch(s.Ctx, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
mockDataProxy.EXPECT().CreateUploadLocation(s.Ctx, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)

err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx)
assert.Nil(t, err)
Expand All @@ -82,7 +82,7 @@ func TestRegisterFromFiles(t *testing.T) {
s.MockAdminClient.EXPECT().CreateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().UpdateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
mockDataProxy := s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient)
mockDataProxy.OnCreateUploadLocationMatch(s.Ctx, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
mockDataProxy.EXPECT().CreateUploadLocation(s.Ctx, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)

err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx)
assert.Nil(t, err)
Expand All @@ -105,7 +105,7 @@ func TestRegisterFromFiles(t *testing.T) {
s.MockAdminClient.EXPECT().CreateWorkflow(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().CreateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().UpdateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).EXPECT().CreateUploadLocation(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
err = Register(s.Ctx, args, config.GetConfig(), s.CmdCtx)
assert.Nil(t, err)
})
Expand Down Expand Up @@ -149,7 +149,7 @@ func TestRegisterFromFiles(t *testing.T) {
s.MockAdminClient.EXPECT().CreateTask(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("failed")).Call.Times(1)
s.MockAdminClient.EXPECT().CreateWorkflow(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("failed")).Call.Times(1)
s.MockAdminClient.EXPECT().CreateLaunchPlan(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("failed")).Call.Times(1)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).EXPECT().CreateUploadLocation(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx)
assert.NotNil(t, err)
assert.Equal(t, fmt.Errorf("failed"), err)
Expand All @@ -174,7 +174,7 @@ func TestRegisterFromFiles(t *testing.T) {
s.MockAdminClient.EXPECT().CreateTask(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("failed")).Call.Times(39)
s.MockAdminClient.EXPECT().CreateWorkflow(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("failed")).Call.Times(21)
s.MockAdminClient.EXPECT().CreateLaunchPlan(mock.Anything, mock.Anything).Return(nil, fmt.Errorf("failed")).Call.Times(24)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).EXPECT().CreateUploadLocation(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx)
assert.NotNil(t, err)
assert.Equal(t, fmt.Errorf("failed"), err)
Expand All @@ -199,7 +199,7 @@ func TestRegisterFromFiles(t *testing.T) {
s.MockAdminClient.EXPECT().CreateWorkflow(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().CreateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().UpdateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).EXPECT().CreateUploadLocation(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx)
assert.Nil(t, err)
})
Expand All @@ -222,7 +222,7 @@ func TestRegisterFromFiles(t *testing.T) {
s.MockAdminClient.EXPECT().CreateTask(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().CreateWorkflow(mock.Anything, mock.Anything).Return(nil, nil)
s.MockAdminClient.EXPECT().CreateLaunchPlan(mock.Anything, mock.Anything).Return(nil, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).EXPECT().CreateUploadLocation(mock.Anything, mock.Anything).Return(&service.CreateUploadLocationResponse{}, nil)
err = registerFromFilesFunc(s.Ctx, args, s.CmdCtx)
assert.Nil(t, err)
})
Expand Down
4 changes: 2 additions & 2 deletions flytectl/cmd/register/register_util_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -403,7 +403,7 @@ func TestUploadFastRegisterArtifact(t *testing.T) {
}, testScope.NewSubScope("flytectl"))
assert.Nil(t, err)
Client = store
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(s.Ctx, &service.CreateUploadLocationRequest{
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).EXPECT().CreateUploadLocation(s.Ctx, &service.CreateUploadLocationRequest{
Project: "flytesnacks",
Domain: "development",
Filename: "flytesnacks-core.tgz",
Expand All @@ -422,7 +422,7 @@ func TestUploadFastRegisterArtifact(t *testing.T) {
}, testScope.NewSubScope("flytectl"))
assert.Nil(t, err)
Client = store
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).OnCreateUploadLocationMatch(s.Ctx, &service.CreateUploadLocationRequest{
s.MockClient.DataProxyClient().(*mocks.DataProxyServiceClient).EXPECT().CreateUploadLocation(s.Ctx, &service.CreateUploadLocationRequest{
Project: "flytesnacks",
Domain: "development",
Filename: "flytesnacks-core.tgz",
Expand Down
10 changes: 5 additions & 5 deletions flyteplugins/go/tasks/plugins/webapi/agent/integration_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -305,19 +305,19 @@ func newMockSyncAgentPlugin() webapi.PluginEntry {
resource := &admin.Resource{Phase: flyteIdlCore.TaskExecution_SUCCEEDED, Outputs: output}

stream := new(agentMocks.SyncAgentService_ExecuteTaskSyncClient)
stream.OnRecv().Return(&admin.ExecuteTaskSyncResponse{
stream.EXPECT().Recv().Return(&admin.ExecuteTaskSyncResponse{
Res: &admin.ExecuteTaskSyncResponse_Header{
Header: &admin.ExecuteTaskSyncResponseHeader{
Resource: resource,
},
},
}, nil).Once()

stream.OnRecv().Return(nil, io.EOF).Once()
stream.OnSendMatch(mock.Anything).Return(nil)
stream.OnCloseSendMatch(mock.Anything).Return(nil)
stream.EXPECT().Recv().Return(nil, io.EOF).Once()
stream.EXPECT().Send(mock.Anything).Return(nil)
stream.EXPECT().CloseSend().Return(nil)

syncAgentClient.OnExecuteTaskSyncMatch(mock.Anything).Return(stream, nil)
syncAgentClient.EXPECT().ExecuteTaskSync(mock.Anything).Return(stream, nil)

cfg := defaultConfig
cfg.DefaultAgent.Endpoint = defaultAgentEndpoint
Expand Down

0 comments on commit 098c2b1

Please sign in to comment.