1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
|
package mock_modserver
import (
"github.com/golang/mock/gomock"
"gitlab.com/gitlab-org/cluster-integration/gitlab-agent/v16/internal/tool/retry"
)
func NewMockAgentRpcApiWithMockPoller(ctrl *gomock.Controller, pollTimes int) *MockAgentRpcApi {
mockRpcApi := NewMockAgentRpcApi(ctrl)
if pollTimes > 0 {
mockRpcApi.EXPECT().
PollWithBackoff(gomock.Any(), gomock.Any()).
DoAndReturn(func(cfg retry.PollConfig, f retry.PollWithBackoffFunc) error {
for i := 0; i < pollTimes; i++ {
err, res := f()
if err != nil || res == retry.Done {
return err
}
}
return nil
})
}
return mockRpcApi
}
|
