test: add missing mocks for SetToken

Author: tronghn

internal/deviceagent/state_machine_test.go

@@ -35,6 +35,7 @@ func TestStateMachine(t *testing.T) {
 		rc.EXPECT().SetAPIServerInfo(mock.Anything, mock.Anything).Return()
 		rc.EXPECT().BuildHelperConfiguration(mock.Anything).Return(&pb.Configuration{})
 		rc.EXPECT().SetJitaToken(mock.Anything).Return()
+		rc.EXPECT().SetToken(mock.Anything).Return()
 
 		mockGetDeviceConfigclient := pb.NewMockAPIServer_GetDeviceConfigurationClient(t)
 		recv := mockGetDeviceConfigclient.EXPECT().Recv().After(20 * time.Millisecond)

internal/deviceagent/states/disconnected/disconnected_test.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"testing"
 
+	"github.com/nais/device/internal/deviceagent/auth"
 	"github.com/nais/device/internal/deviceagent/config"
 	"github.com/nais/device/internal/deviceagent/runtimeconfig"
 	statemachine "github.com/nais/device/internal/deviceagent/statemachine/state"
@@ -19,6 +20,7 @@ func TestDisconnected(t *testing.T) {
 		rc.EXPECT().ResetEnrollConfig()
 		rc.EXPECT().SetAPIServerInfo(nil, "").Return()
 		rc.EXPECT().SetJitaToken((*oauth2.Token)(nil)).Return()
+		rc.EXPECT().SetToken((*auth.Tokens)(nil)).Return()
 
 		cfg := config.Config{
 			AgentConfiguration: &pb.AgentConfiguration{
@@ -38,6 +40,7 @@ func TestDisconnected(t *testing.T) {
 		rc.EXPECT().ResetEnrollConfig()
 		rc.EXPECT().SetAPIServerInfo(nil, "").Return()
 		rc.EXPECT().SetJitaToken((*oauth2.Token)(nil)).Return()
+		rc.EXPECT().SetToken((*auth.Tokens)(nil)).Return()
 
 		cfg := config.Config{
 			AgentConfiguration: &pb.AgentConfiguration{

internal/integration_test/integration_test.go

@@ -253,6 +253,7 @@ func tableTest(t *testing.T, log *logrus.Entry, testDevice *pb.Device, endState
 	rc.EXPECT().APIServerPeer().Return(apiserverPeer)
 	rc.EXPECT().SetAPIServerInfo(mock.Anything, mock.Anything).Return()
 	rc.EXPECT().SetJitaToken(mock.Anything).Return().Maybe()
+	rc.EXPECT().SetToken(mock.Anything).Return().Maybe()
 
 	initialHelperConfig := &pb.Configuration{
 		Gateways:   mapValues(initialPeers),