// Copyright 2021 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"errors"
	"fmt"
	"reflect"
	"testing"

	"policy_manager/dbus"
	"policy_manager/mock/mockdbus"
	"policy_manager/mock/mockdevicepolicy"
	"policy_manager/mock/mocksysapi"
	"policy_manager/policymanagerproto"

	"github.com/golang/mock/gomock"
	"github.com/golang/protobuf/proto"
)

// Dummy test data.
var devChannel = policymanagerproto.ReleaseChannel_DEV
var betaChannel = policymanagerproto.ReleaseChannel_BETA
var stableChannel = policymanagerproto.ReleaseChannel_STABLE

var dummyChannel = betaChannel
var dummyStatus = &policymanagerproto.InstanceStatus{
	OsVersion: &policymanagerproto.OSVersion{
		VersionString: proto.String("16108.403.11"),
		Milestone:     proto.Uint32(89),
		Channel:       &dummyChannel,
	},
}

var emptyUserConfig = &policymanagerproto.InstanceConfig{}

var dummyUserConfig = &policymanagerproto.InstanceConfig{
	MetricsEnabled:      proto.Bool(true),
	TargetVersionPrefix: proto.String("7121."),
	RebootAfterUpdate:   proto.Bool(false),
}

var dummyInstanceConfig = &policymanagerproto.InstanceConfig{
	MetricsEnabled:      proto.Bool(true),
	TargetVersionPrefix: proto.String("7121."),
	RebootAfterUpdate:   proto.Bool(false),
}

// TestUpdateInstanceConfig tests that we are making the expected calls to fetch
// instance config and generate new device policy.
func TestUpdateInstanceConfig(t *testing.T) {
	tests := []struct {
		name                        string
		userConfig                  *policymanagerproto.InstanceConfig
		status                      *policymanagerproto.InstanceStatus
		onDiskConfig                *policymanagerproto.InstanceConfig
		managerGetInstanceConfigErr error
		expectChannelSwitch         string
		setChannelErr               error
		expectedUserConfig          *policymanagerproto.InstanceConfig
		setInstanceConfigErr        error
		expectErr                   bool
	}{
		{
			name:                        "SuccessfulUpdateConfigFetch",
			userConfig:                  emptyUserConfig,
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name:                        "NilInstanceStatus",
			userConfig:                  emptyUserConfig,
			status:                      nil,
			onDiskConfig:                emptyUserConfig,
			managerGetInstanceConfigErr: nil,
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig:          nil,
			setInstanceConfigErr:        nil,
			expectErr:                   true,
		},
		{
			name:                        "NilUserConfig",
			userConfig:                  nil,
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig:          nil,
			setInstanceConfigErr:        nil,
			expectErr:                   false,
		},
		{
			name:                        "LocallyResolvedInstanceConfig",
			userConfig:                  dummyUserConfig,
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				MetricsEnabled:      proto.Bool(true),
				TargetVersionPrefix: proto.String("7121."),
				RebootAfterUpdate:   proto.Bool(false),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name: "UpdateDisabled",
			userConfig: &policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("update_disabled"),
			},
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				UpdateStrategy:      proto.String("update_disabled"),
				TargetVersionPrefix: proto.String("16108.403.11"),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name: "ImageNameWithoutChannelSwitch",
			userConfig: &policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("cos-dev-93-16511-0-0"),
			},
			status: &policymanagerproto.InstanceStatus{
				OsVersion: &policymanagerproto.OSVersion{
					VersionString: proto.String("16511.0.0"),
					Milestone:     proto.Uint32(93),
					Channel:       &devChannel,
				},
			},
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				UpdateStrategy:      proto.String("cos-dev-93-16511-0-0"),
				TargetVersionPrefix: proto.String("16511.0.0"),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name: "ImageNameWithChannelSwitch",
			userConfig: &policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("cos-dev-93-16511-0-0"),
			},
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "dev",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				UpdateStrategy:      proto.String("cos-dev-93-16511-0-0"),
				TargetVersionPrefix: proto.String("16511.0.0"),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name: "ChannelSwitchWithErr",
			userConfig: &policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("cos-dev-93-16511-0-0"),
			},
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "dev",
			setChannelErr:               errors.New("update engine failed to set channel"),
			expectedUserConfig:          nil,
			setInstanceConfigErr:        nil,
			expectErr:                   true,
		},
		{
			name: "MetricsOnlyUserConfig",
			userConfig: &policymanagerproto.InstanceConfig{
				MetricsEnabled: proto.Bool(true),
			},
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				MetricsEnabled: proto.Bool(true),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name:                        "NoOnDiskConfig",
			userConfig:                  emptyUserConfig,
			status:                      dummyStatus,
			onDiskConfig:                nil,
			managerGetInstanceConfigErr: errors.New("missing policy file"),
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name:       "EnforcedOnDiskConfig",
			userConfig: dummyUserConfig,
			status:     dummyStatus,
			onDiskConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					LoggingEnabled:    proto.Bool(true),
					MonitoringEnabled: proto.Bool(false),
				},
			},
			managerGetInstanceConfigErr: nil,
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				MetricsEnabled:      proto.Bool(true),
				TargetVersionPrefix: proto.String("7121."),
				RebootAfterUpdate:   proto.Bool(false),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(true),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name: "EnforcedUserConfig",
			userConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					MonitoringEnabled: proto.Bool(true),
				},
			},
			status:                      dummyStatus,
			onDiskConfig:                emptyUserConfig,
			managerGetInstanceConfigErr: nil,
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					MonitoringEnabled: proto.Bool(true),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
		{
			name:       "NotEnforceNilUserConfig",
			userConfig: nil,
			status:     dummyStatus,
			onDiskConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					LoggingEnabled:    proto.Bool(true),
					MonitoringEnabled: proto.Bool(false),
				},
			},
			managerGetInstanceConfigErr: nil,
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig:          nil,
			setInstanceConfigErr:        nil,
			expectErr:                   false,
		},
		{
			name:       "NoEnforcement",
			userConfig: dummyUserConfig,
			status:     dummyStatus,
			onDiskConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			managerGetInstanceConfigErr: nil,
			expectChannelSwitch:         "",
			setChannelErr:               nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				MetricsEnabled:      proto.Bool(true),
				TargetVersionPrefix: proto.String("7121."),
				RebootAfterUpdate:   proto.Bool(false),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
			setInstanceConfigErr: nil,
			expectErr:            false,
		},
	}
	for _, test := range tests {
		mockCtrl := gomock.NewController(t)
		mockManager := mockdevicepolicy.NewMockManager(mockCtrl)
		mockUEClient := mockdbus.NewMockUpdateEngineClient(mockCtrl)
		userConfig := proto.Clone(test.userConfig).(*policymanagerproto.InstanceConfig)

		t.Logf("Running: %s", test.name)

		// Reading on-disk policy will return onDiskConfig before the file is overwritten.
		readPolicyFile := mockManager.EXPECT().GetInstanceConfig().Return(test.onDiskConfig, test.managerGetInstanceConfigErr).AnyTimes()

		if test.expectChannelSwitch != "" {
			channelConfig := fmt.Sprintf("%s-channel", test.expectChannelSwitch)
			mockUEClient.EXPECT().SetChannel(channelConfig).Return(test.setChannelErr)
		}

		// Only when status is provided, and ueClient can correctly switch channel,
		// updateInstanceConfig() will succeed.
		if test.status != nil && test.setChannelErr == nil {
			mockManager.EXPECT().SetInstanceConfig(test.expectedUserConfig).Return(test.setInstanceConfigErr).After(readPolicyFile)
		}

		err := updateInstanceConfig(test.status, userConfig, mockManager, mockUEClient)
		if err != nil && !test.expectErr {
			t.Errorf("test %s got unexpected error: %v", test.name,
				err)
		} else if err == nil && test.expectErr {
			t.Errorf("test %s passed, expected error", test.name)
		}

		mockCtrl.Finish()
	}
}

func TestResolveEnforcementConfig(t *testing.T) {
	tests := []struct {
		name                 string
		userConfig           *policymanagerproto.InstanceConfig
		onDiskConfig         *policymanagerproto.InstanceConfig
		getInstanceConfigErr error
		expectedUserConfig   *policymanagerproto.InstanceConfig
	}{
		{
			name:                 "NoOnDiskConfig",
			userConfig:           emptyUserConfig,
			onDiskConfig:         nil,
			getInstanceConfigErr: errors.New("missing policy file"),
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
		},
		{
			name:       "EnforcedOnDiskConfig",
			userConfig: dummyUserConfig,
			onDiskConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					LoggingEnabled:    proto.Bool(true),
					MonitoringEnabled: proto.Bool(false),
				},
			},
			getInstanceConfigErr: nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				MetricsEnabled:      proto.Bool(true),
				TargetVersionPrefix: proto.String("7121."),
				RebootAfterUpdate:   proto.Bool(false),
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(true),
				},
			},
		},
		{
			name: "EnforcedUserConfig",
			userConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					MonitoringEnabled: proto.Bool(true),
				},
			},
			onDiskConfig:         emptyUserConfig,
			getInstanceConfigErr: nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					MonitoringEnabled: proto.Bool(true),
				},
			},
		},
		{
			name:       "NilUserConfig",
			userConfig: nil,
			onDiskConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced:          proto.Bool(true),
					LoggingEnabled:    proto.Bool(true),
					MonitoringEnabled: proto.Bool(false),
				},
			},
			getInstanceConfigErr: nil,
			expectedUserConfig:   nil,
		},
		{
			name:                 "NoEnforcement",
			userConfig:           emptyUserConfig,
			onDiskConfig:         dummyUserConfig,
			getInstanceConfigErr: nil,
			expectedUserConfig: &policymanagerproto.InstanceConfig{
				HealthMonitorConfig: &policymanagerproto.HealthMonitorConfig{
					Enforced: proto.Bool(false),
				},
			},
		},
	}

	for _, test := range tests {
		mockCtrl := gomock.NewController(t)
		mockManager := mockdevicepolicy.NewMockManager(mockCtrl)
		userConfig := proto.Clone(test.userConfig).(*policymanagerproto.InstanceConfig)

		t.Logf("Running: %s", test.name)

		mockManager.EXPECT().GetInstanceConfig().Return(test.onDiskConfig, test.getInstanceConfigErr).AnyTimes()

		resolveEnforcementConfig(userConfig, mockManager)

		if !proto.Equal(userConfig, test.expectedUserConfig) {
			t.Errorf("In test %s, got %s, expected %s",
				test.name, userConfig, test.expectedUserConfig)
		}

		mockCtrl.Finish()
	}
}

func TestResolveLocalUpdateStrategy(t *testing.T) {
	dummyOsVersion := &policymanagerproto.OSVersion{
		VersionString: proto.String("16108.403.11"),
		Milestone:     proto.Uint32(89),
		Channel:       &dummyChannel,
	}
	tests := []struct {
		name               string
		userConfig         *policymanagerproto.InstanceConfig
		osversion          *policymanagerproto.OSVersion
		resolvedUserConfig *policymanagerproto.InstanceConfig
		result             error
	}{
		{
			"NilUserConfig",
			nil,
			dummyOsVersion,
			nil,
			nil,
		},
		{
			"NilOsVersion",
			dummyInstanceConfig,
			nil,
			dummyInstanceConfig,
			errors.New("osversion is nil, which should never happen"),
		},
		{
			"EmptyUserConfig",
			emptyUserConfig,
			dummyOsVersion,
			emptyUserConfig,
			nil,
		},
		{
			"SpecifcTargetVersionPrefix",
			&policymanagerproto.InstanceConfig{
				TargetVersionPrefix: proto.String("1.2.3"),
			},
			dummyOsVersion,
			&policymanagerproto.InstanceConfig{
				TargetVersionPrefix: proto.String("1.2.3"),
			},
			nil,
		},
		{
			"EmptyUpdateStrategy",
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String(""),
			},
			dummyOsVersion,
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String(""),
			},
			nil,
		},
		{
			"UpdateDisabledUpdateStrategy",
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("update_disabled"),
			},
			dummyOsVersion,
			&policymanagerproto.InstanceConfig{
				UpdateStrategy:      proto.String("update_disabled"),
				TargetVersionPrefix: proto.String("16108.403.11"),
			},
			nil,
		},
		{
			"FullImageNameUpdateStrategyCOS",
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("cos-beta-89-16108-403-11"),
			},
			dummyOsVersion,
			&policymanagerproto.InstanceConfig{
				UpdateStrategy:      proto.String("cos-beta-89-16108-403-11"),
				TargetVersionPrefix: proto.String("16108.403.11"),
			},
			nil,
		},
		{
			"GarbageUpdateStrategy",
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("a-b-c-d-12-345-67-89"),
			},
			dummyOsVersion,
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("a-b-c-d-12-345-67-89"),
			},
			errors.New("user hasn't specified any strategy we know"),
		},
		{
			"BadChannelUpdateStrategy",
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("cos-foo-89-16108-403-11"),
			},
			dummyOsVersion,
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("cos-foo-89-16108-403-11"),
			},
			errors.New("user hasn't specified any strategy we know"),
		},
		{
			"UnknownUpdateStrategy",
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("stable-89"),
			},
			dummyOsVersion,
			&policymanagerproto.InstanceConfig{
				UpdateStrategy: proto.String("stable-89"),
			},
			errors.New("user hasn't specified any strategy we know"),
		},
	}

	for _, test := range tests {
		userConfig := proto.Clone(test.userConfig).(*policymanagerproto.InstanceConfig)

		ret := resolveLocalUpdateStrategy(userConfig, test.osversion)
		if ret != test.result && ret.Error() != test.result.Error() {
			t.Errorf("test '%s' failed. Unexpected return value '%v'; Expected '%v'", test.name, ret, test.result)
		}
		if !proto.Equal(userConfig, test.resolvedUserConfig) {
			t.Errorf("test '%s' failed. Unexpected resolved config: %v; Expected: %v", test.name, userConfig, test.resolvedUserConfig)
		}
	}
}

// This tests getOSVersionFromStrategy parses image name and returns
// os version correctly.
func TestGetOSVersionFromStrategy(t *testing.T) {
	testCases := []struct {
		name            string
		strategy        string
		expectedVersion *policymanagerproto.OSVersion
		expectErr       bool
	}{
		{
			"EmptyStrategy",
			"",
			nil,
			true,
		},
		{
			"PositiveCaseCosDev",
			"cos-dev-93-16511-0-0",
			&policymanagerproto.OSVersion{
				VersionString: proto.String("16511.0.0"),
				Milestone:     proto.Uint32(uint32(93)),
				Channel:       &devChannel,
			},
			false,
		},
		{
			"PositiveCaseCosBeta",
			"cos-beta-89-16108-403-11",
			&policymanagerproto.OSVersion{
				VersionString: proto.String("16108.403.11"),
				Milestone:     proto.Uint32(uint32(89)),
				Channel:       &betaChannel,
			},
			false,
		},
		{
			"PositiveCaseCosStable",
			"cos-stable-89-16108-470-1",
			&policymanagerproto.OSVersion{
				VersionString: proto.String("16108.470.1"),
				Milestone:     proto.Uint32(uint32(89)),
				Channel:       &stableChannel,
			},
			false,
		},
		{
			"NegativeCaseNoCodeName",
			"stable-89-16108-470-1",
			nil,
			true,
		},
		{
			"NegativeCaseNoChannelName",
			"cos-89-16108-470-1",
			nil,
			true,
		},
		{
			"NegativeCaseBadChannel",
			"cos-foo-89-16108-470-1",
			nil,
			true,
		},
		{
			"NegativeCaseLessDigit",
			"cos-stable-89-16108-470",
			nil,
			true,
		},
		{
			"NegativeCaseMoreDigit",
			"cos-stable-89-16108-470-1-0",
			nil,
			true,
		},
	}
	for _, test := range testCases {
		actualVersion, err := getOSVersionFromStrategy(test.strategy)
		if !reflect.DeepEqual(actualVersion, test.expectedVersion) {
			t.Errorf("test %s: expected image version: %v, got: %v", test.name, test.expectedVersion, actualVersion)
		}
		if err != nil && !test.expectErr {
			t.Errorf("test %s got unexpected error: %v", test.name, err)
		} else if err == nil && test.expectErr {
			t.Errorf("test %s passed, expected error", test.name)
		}
	}
}

// This tests processUEStatus processes ueStatus correctly.
func TestProcessUEStatus(t *testing.T) {
	testCases := []struct {
		name           string
		ueStatus       dbus.UEGetStatusResponse
		instanceConfig *policymanagerproto.InstanceConfig
		expectReboot   bool
		expectErr      bool
		getConfigError error
	}{
		{
			"IdleStatus",
			dbus.UEGetStatusResponse{
				UpdateStatus: "UPDATE_STATUS_IDLE",
			},
			nil,
			false,
			false,
			nil,
		},
		{
			"NeedReboot",
			dbus.UEGetStatusResponse{
				UpdateStatus: "UPDATE_STATUS_UPDATED_NEED_REBOOT",
				NewVersion:   "12345.0.0",
			},
			&policymanagerproto.InstanceConfig{
				TargetVersionPrefix: proto.String("12345.0.0"),
				RebootAfterUpdate:   proto.Bool(true),
			},
			true,
			false,
			nil,
		},
		{
			"NeedRebootNoTargetVersion",
			dbus.UEGetStatusResponse{
				UpdateStatus: "UPDATE_STATUS_UPDATED_NEED_REBOOT",
				NewVersion:   "12345.0.0",
			},
			&policymanagerproto.InstanceConfig{
				TargetVersionPrefix: proto.String(""),
				RebootAfterUpdate:   proto.Bool(true),
			},
			true,
			false,
			nil,
		},

		{
			"VersionMismatch",
			dbus.UEGetStatusResponse{
				UpdateStatus: "UPDATE_STATUS_UPDATED_NEED_REBOOT",
				NewVersion:   "12345.0.0",
			},
			&policymanagerproto.InstanceConfig{
				TargetVersionPrefix: proto.String("54321.0.0"),
				RebootAfterUpdate:   proto.Bool(true),
			},
			false,
			true,
			nil,
		},
		{
			"UserDontNeedReboot",
			dbus.UEGetStatusResponse{
				UpdateStatus: "UPDATE_STATUS_UPDATED_NEED_REBOOT",
				NewVersion:   "12345.0.0",
			},
			&policymanagerproto.InstanceConfig{
				TargetVersionPrefix: proto.String("12345.0.0"),
				RebootAfterUpdate:   proto.Bool(false),
			},
			false,
			false,
			nil,
		},
		{
			"GetError",
			dbus.UEGetStatusResponse{
				UpdateStatus: "UPDATE_STATUS_UPDATED_NEED_REBOOT",
				NewVersion:   "12345.0.0",
			},
			nil,
			false,
			true,
			fmt.Errorf("GetInstanceConfig Error"),
		},
	}
	for _, test := range testCases {
		mockCtrl := gomock.NewController(t)
		mockManager := mockdevicepolicy.NewMockManager(mockCtrl)
		mockAPIHandler := mocksysapi.NewMockAPIHandler(mockCtrl)
		mockManager.EXPECT().GetInstanceConfig().Return(
			test.instanceConfig, test.getConfigError)
		if test.expectReboot {
			mockAPIHandler.EXPECT().RunCommand("/sbin/shutdown", "-r", "now").Return(
				nil, nil, nil)
		}

		err := processUEStatus(test.ueStatus, mockManager, mockAPIHandler)

		if err == nil && test.expectErr {
			t.Errorf("test %s passed, expected error", test.name)
		} else if err != nil && !test.expectErr {
			t.Errorf("test %s got unexpected error: %v", test.name, err)
		}
	}
}