package mock_modserver

import (
	"gitlab.com/gitlab-org/cluster-integration/tunnel/tool/retry"
	"go.uber.org/mock/gomock"
)

func NewMockAgentRPCAPIWithMockPoller(ctrl *gomock.Controller, pollTimes int) *MockAgentRPCAPI {
	mockRPCAPI := NewMockAgentRPCAPI(ctrl)
	if pollTimes > 0 {
		mockRPCAPI.EXPECT().
			PollWithBackoff(gomock.Any(), gomock.Any(), gomock.Any()).
			DoAndReturn(func(_ retry.PollConfig, f retry.PollWithBackoffFunc, opts ...retry.PollOption) error {
				for range pollTimes {
					err, res := f()
					if err != nil || res == retry.Done {
						return err
					}
				}
				return nil
			})
	}
	return mockRPCAPI
}
