/*
 * Copyright (C) 2022 The Android Open Source Project
 *
 * 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 com.android.devicelockcontroller.provision.worker;

import static android.content.pm.PackageManager.COMPONENT_ENABLED_STATE_DEFAULT;
import static android.net.NetworkCapabilities.NET_CAPABILITY_INTERNET;
import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_RESTRICTED;
import static android.net.NetworkCapabilities.NET_CAPABILITY_NOT_VPN;
import static android.net.NetworkCapabilities.NET_CAPABILITY_TRUSTED;
import static android.os.Looper.getMainLooper;

import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceIdType.DEVICE_ID_TYPE_IMEI;
import static com.android.devicelockcontroller.common.DeviceLockConstants.DeviceIdType.DEVICE_ID_TYPE_MEID;
import static com.android.devicelockcontroller.common.DeviceLockConstants.READY_FOR_PROVISION;
import static com.android.devicelockcontroller.common.DeviceLockConstants.RETRY_CHECK_IN;
import static com.android.devicelockcontroller.common.DeviceLockConstants.STOP_CHECK_IN;
import static com.android.devicelockcontroller.provision.worker.GetFcmTokenWorker.FCM_TOKEN_WORKER_INITIAL_DELAY;
import static com.android.devicelockcontroller.provision.worker.GetFcmTokenWorker.FCM_TOKEN_WORK_NAME;
import static com.android.devicelockcontroller.stats.StatsLogger.CheckInRetryReason.CONFIG_UNAVAILABLE;

import static com.google.common.truth.Truth.assertThat;

import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.robolectric.annotation.LooperMode.Mode.LEGACY;

import android.content.Intent;
import android.content.pm.PackageInfo;
import android.content.pm.PackageManager;
import android.net.NetworkRequest;
import android.os.SystemClock;
import android.telephony.TelephonyManager;
import android.util.ArraySet;

import androidx.annotation.Nullable;
import androidx.test.core.app.ApplicationProvider;
import androidx.work.Configuration;
import androidx.work.WorkInfo;
import androidx.work.WorkManager;
import androidx.work.testing.SynchronousExecutor;
import androidx.work.testing.WorkManagerTestInitHelper;

import com.android.devicelockcontroller.FeatureFlagProvider;
import com.android.devicelockcontroller.TestDeviceLockControllerApplication;
import com.android.devicelockcontroller.common.DeviceId;
import com.android.devicelockcontroller.common.DeviceLockConstants.DeviceCheckInStatus;
import com.android.devicelockcontroller.policy.DevicePolicyController;
import com.android.devicelockcontroller.policy.FinalizationController;
import com.android.devicelockcontroller.policy.FinalizationControllerImpl;
import com.android.devicelockcontroller.policy.ProvisionStateController;
import com.android.devicelockcontroller.provision.grpc.GetDeviceCheckInStatusGrpcResponse;
import com.android.devicelockcontroller.provision.grpc.ProvisioningConfiguration;
import com.android.devicelockcontroller.receivers.ProvisionReadyReceiver;
import com.android.devicelockcontroller.schedule.DeviceLockControllerScheduler;
import com.android.devicelockcontroller.stats.StatsLogger;
import com.android.devicelockcontroller.stats.StatsLoggerProvider;
import com.android.devicelockcontroller.storage.GlobalParametersClient;

import com.google.common.util.concurrent.Futures;

import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.mockito.Mockito;
import org.robolectric.RobolectricTestRunner;
import org.robolectric.Shadows;
import org.robolectric.annotation.LooperMode;
import org.robolectric.shadows.ShadowPackageManager;
import org.robolectric.shadows.ShadowTelephonyManager;

import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.Objects;

@LooperMode(LEGACY)
@RunWith(RobolectricTestRunner.class)
public final class DeviceCheckInHelperTest {
    static final Duration TEST_CHECK_RETRY_DURATION = Duration.ofDays(30);
    static final Duration TEST_NEGATIVE_CHECK_RETRY_DURATION =
            Duration.ZERO.minus(TEST_CHECK_RETRY_DURATION);
    public static final boolean IS_PROVISIONING_MANDATORY = false;
    public static final String PHONE_SYSTEM_APP_PACKAGE_NAME = "com.android.phone";
    private TestDeviceLockControllerApplication mTestApplication;
    static final int TOTAL_SLOT_COUNT = 2;
    static final int TOTAL_ID_COUNT = 4;
    static final String IMEI_1 = "IMEI1";
    static final String IMEI_2 = "IMEI2";
    static final String MEID_1 = "MEID1";
    static final String MEID_2 = "MEID2";
    static final ArraySet<DeviceId> ACTUAL_DEVICE_IDs =
            new ArraySet<>(
                    new DeviceId[] {
                        new DeviceId(DEVICE_ID_TYPE_IMEI, IMEI_1),
                        new DeviceId(DEVICE_ID_TYPE_IMEI, IMEI_2),
                        new DeviceId(DEVICE_ID_TYPE_MEID, MEID_1),
                        new DeviceId(DEVICE_ID_TYPE_MEID, MEID_2),
                    });
    static final ProvisioningConfiguration TEST_CONFIGURATION =
            new ProvisioningConfiguration(
                    /* kioskAppProviderName= */ "test_provider",
                    /* kioskAppPackageName= */ "test_package",
                    /* kioskAppAllowlistPackages= */ List.of(
                            "test_allowed_app1", "test_allowed_app2"),
                    /* kioskAppEnableOutgoingCalls= */ false,
                    /* kioskAppEnableEnableNotifications= */ true,
                    /* disallowInstallingFromUnknownSources= */ false,
                    /* termsAndConditionsUrl= */ "test_terms_and_configurations_url",
                    /* supportUrl= */ "test_support_url");
    static final int DEVICE_ID_TYPE_BITMAP =
            (1 << DEVICE_ID_TYPE_IMEI) | (1 << DEVICE_ID_TYPE_MEID);
    static final String FAKE_APEX_PACKAGE = "fake_apex";
    static final long APEX_VERSION = 1111;
    static final String DEVICE_LOCALE = "en-US";

    private FinalizationController mFinalizationController;
    private DeviceCheckInHelper mHelper;

    private ShadowTelephonyManager mTelephonyManager;
    private GlobalParametersClient mGlobalParametersClient;
    private DeviceLockControllerScheduler mScheduler;
    private StatsLogger mStatsLogger;
    private WorkManager mWorkManager;
    private PackageManager mPackageManager;
    private ShadowPackageManager mShadowPackageManager;
    private FeatureFlagProvider mFeatureFlagProvider;
    private DevicePolicyController mMockDevicePolicyController;
    private ProvisionStateController mMockProvisionStateController;

    @Before
    public void setUp() {
        mTestApplication = ApplicationProvider.getApplicationContext();
        mScheduler = mTestApplication.getDeviceLockControllerScheduler();
        mFinalizationController = mTestApplication.getFinalizationController();
        when(mFinalizationController.notifyRestrictionsCleared())
                .thenReturn(Futures.immediateVoidFuture());
        when(mFinalizationController.finalizeNotEnrolledDevice())
                .thenReturn(Futures.immediateVoidFuture());
        when(mFinalizationController.enforceDiskState(true))
                .thenReturn(Futures.immediateVoidFuture());

        mFeatureFlagProvider = mTestApplication.getFeatureFlagProvider();
        mTelephonyManager =
                Shadows.shadowOf(mTestApplication.getSystemService(TelephonyManager.class));
        mHelper = new DeviceCheckInHelper(mTestApplication);
        WorkManagerTestInitHelper.initializeTestWorkManager(
                mTestApplication,
                new Configuration.Builder()
                        .setMinimumLoggingLevel(android.util.Log.DEBUG)
                        .setExecutor(new SynchronousExecutor())
                        .build());
        mWorkManager = WorkManager.getInstance(mTestApplication);
        mGlobalParametersClient = GlobalParametersClient.getInstance();
        mStatsLogger = ((StatsLoggerProvider) mTestApplication).getStatsLogger();
        mPackageManager = mTestApplication.getPackageManager();
        mShadowPackageManager = Shadows.shadowOf(mPackageManager);
        mMockDevicePolicyController = mTestApplication.getPolicyController();
        mMockProvisionStateController = mTestApplication.getProvisionStateController();
        when(mMockProvisionStateController.notifyProvisioningReady())
                .thenReturn(Futures.immediateVoidFuture());
    }

    @Test
    public void getDeviceAvailableUniqueIds_shouldReturnAllAvailableUniqueIds() {
        mShadowPackageManager.setSystemFeature(
                PackageManager.FEATURE_TELEPHONY_CDMA, /* supported= */ true);
        mTelephonyManager.setActiveModemCount(TOTAL_SLOT_COUNT);
        mTelephonyManager.setImei(/* slotIndex= */ 0, IMEI_1);
        mTelephonyManager.setImei(/* slotIndex= */ 1, IMEI_2);
        mTelephonyManager.setMeid(/* slotIndex= */ 0, MEID_1);
        mTelephonyManager.setMeid(/* slotIndex= */ 1, MEID_2);
        final ArraySet<DeviceId> deviceIds =
                mHelper.getDeviceAvailableUniqueIds(DEVICE_ID_TYPE_BITMAP);
        assertThat(Objects.requireNonNull(deviceIds).size()).isEqualTo(TOTAL_ID_COUNT);
        assertThat(deviceIds).containsExactlyElementsIn(ACTUAL_DEVICE_IDs);
    }

    @Test
    public void handleGetDeviceCheckInStatusResponse_stopCheckIn_finalizesNonEnrolledDevice()
            throws Exception {
        final GetDeviceCheckInStatusGrpcResponse response = createStopResponse();

        assertThat(
                        mHelper.handleGetDeviceCheckInStatusResponse(
                                response,
                                mock(DeviceLockControllerScheduler.class),
                                mTestApplication.getFcmRegistrationToken().get()))
                .isTrue();
        Shadows.shadowOf(getMainLooper()).idle();
        verify(mFinalizationController).finalizeNotEnrolledDevice();
    }

    @Test
    public void handleProvisionReadyResponse_validConfiguration_shouldSendBroadcast() {
        GetDeviceCheckInStatusGrpcResponse response = createReadyResponse(TEST_CONFIGURATION);

        assertThat(mHelper.handleProvisionReadyResponse(response)).isTrue();

        assertThat(Futures.getUnchecked(mGlobalParametersClient.isProvisionReady())).isTrue();
        List<Intent> intents = Shadows.shadowOf(mTestApplication).getBroadcastIntents();
        assertThat(intents.size()).isEqualTo(1);
        assertThat(intents.get(0).getComponent().getClassName())
                .isEqualTo(ProvisionReadyReceiver.class.getName());
    }

    @Test
    public void handleProvisionReadyResponse_recolEnabled_shouldChangeFinalizedToUnfinalized() {
        when(mFeatureFlagProvider.isRecolEnabled()).thenReturn(true);
        Futures.getUnchecked(
                GlobalParametersClient.getInstance()
                        .setFinalizationState(
                                FinalizationControllerImpl.FinalizationState.FINALIZED));
        GetDeviceCheckInStatusGrpcResponse response = createReadyResponse(TEST_CONFIGURATION);

        assertThat(mHelper.handleProvisionReadyResponse(response)).isTrue();

        assertThat(
                        Futures.getUnchecked(
                                GlobalParametersClient.getInstance().getFinalizationState()))
                .isEqualTo(FinalizationControllerImpl.FinalizationState.UNFINALIZED);
        verify(mFinalizationController).enforceDiskState(/* force= */ true);
    }

    @Test
    public void handleProvisionReadyResponse_recolNotEnabled_shouldNotChangeFinalizedState() {
        when(mFeatureFlagProvider.isRecolEnabled()).thenReturn(false);
        Futures.getUnchecked(
                GlobalParametersClient.getInstance()
                        .setFinalizationState(
                                FinalizationControllerImpl.FinalizationState.FINALIZED));
        GetDeviceCheckInStatusGrpcResponse response = createReadyResponse(TEST_CONFIGURATION);

        assertThat(mHelper.handleProvisionReadyResponse(response)).isTrue();

        assertThat(
                        Futures.getUnchecked(
                                GlobalParametersClient.getInstance().getFinalizationState()))
                .isEqualTo(FinalizationControllerImpl.FinalizationState.FINALIZED);
        verify(mFinalizationController, never()).enforceDiskState(anyBoolean());
    }

    @Test
    public void handleProvisionReadyResponse_invalidConfiguration_shouldNotSendBroadcast() {
        GetDeviceCheckInStatusGrpcResponse response =
                createReadyResponse(/* configuration= */ null);

        assertThat(mHelper.handleProvisionReadyResponse(response)).isFalse();

        assertThat(Futures.getUnchecked(mGlobalParametersClient.isProvisionReady())).isFalse();
        List<Intent> intents = Shadows.shadowOf(mTestApplication).getBroadcastIntents();
        assertThat(intents.size()).isEqualTo(0);
    }

    @Test
    public void handleProvisionReadyResponse_invalidConfiguration_shouldLogRetryCheckIn() {
        GetDeviceCheckInStatusGrpcResponse response =
                createReadyResponse(/* configuration= */ null);

        mHelper.handleProvisionReadyResponse(response);

        verify(mStatsLogger).logCheckInRetry(CONFIG_UNAVAILABLE);
    }

    @Test
    public void handleGetDeviceCheckInStatusResponse_ready_nonEmptyFcmDoesNotStartFcmWork()
            throws Exception {
        final GetDeviceCheckInStatusGrpcResponse response = createReadyResponse(TEST_CONFIGURATION);

        assertThat(
                        mHelper.handleGetDeviceCheckInStatusResponse(
                                response,
                                mScheduler,
                                mTestApplication.getFcmRegistrationToken().get()))
                .isTrue();

        assertThat(mWorkManager.getWorkInfosForUniqueWork(FCM_TOKEN_WORK_NAME).get()).isEmpty();
    }

    @Test
    public void handleGetDeviceCheckInStatusResponse_readyForProvisioning_emptyFcmStartsWork()
            throws Exception {
        final GetDeviceCheckInStatusGrpcResponse response = createReadyResponse(TEST_CONFIGURATION);

        assertThat(
                        mHelper.handleGetDeviceCheckInStatusResponse(
                                response, mScheduler, /* fcmRegistrationToken= */ null))
                .isTrue();

        List<WorkInfo> actualWorks =
                Futures.getUnchecked(mWorkManager.getWorkInfosForUniqueWork(FCM_TOKEN_WORK_NAME));
        assertThat(actualWorks.size()).isEqualTo(1);
        WorkInfo actualWorkInfo = actualWorks.get(0);

        NetworkRequest networkRequest = actualWorkInfo.getConstraints().getRequiredNetworkRequest();
        assertNetworkRequestCapabilities(networkRequest);
        assertThat(actualWorkInfo.getInitialDelayMillis())
                .isEqualTo(FCM_TOKEN_WORKER_INITIAL_DELAY.toMillis());
    }

    @Test
    public void handleGetDeviceCheckInStatusResponse_retryCheckIn_shouldScheduleRetryWork()
            throws Exception {
        final GetDeviceCheckInStatusGrpcResponse response =
                createRetryResponse(
                        SystemClock.currentNetworkTimeClock()
                                .instant()
                                .plus(TEST_CHECK_RETRY_DURATION));

        assertThat(
                        mHelper.handleGetDeviceCheckInStatusResponse(
                                response,
                                mScheduler,
                                mTestApplication.getFcmRegistrationToken().get()))
                .isTrue();

        verify(mScheduler).scheduleRetryCheckInWork(eq(TEST_CHECK_RETRY_DURATION));
    }

    @Test
    public void handleGetDeviceCheckInStatusResponse_retryCheckIn_nonEmptyFcmDoesNotStartFcmWork()
            throws Exception {
        final GetDeviceCheckInStatusGrpcResponse response =
                createRetryResponse(
                        SystemClock.currentNetworkTimeClock()
                                .instant()
                                .plus(TEST_CHECK_RETRY_DURATION));

        assertThat(
                        mHelper.handleGetDeviceCheckInStatusResponse(
                                response,
                                mScheduler,
                                mTestApplication.getFcmRegistrationToken().get()))
                .isTrue();

        assertThat(mWorkManager.getWorkInfosForUniqueWork(FCM_TOKEN_WORK_NAME).get()).isEmpty();
    }

    @Test
    public void handleGetDeviceCheckInStatusResponse_retryCheckIn_emptyFcmStartsWork()
            throws Exception {
        final GetDeviceCheckInStatusGrpcResponse response =
                createRetryResponse(
                        SystemClock.currentNetworkTimeClock()
                                .instant()
                                .plus(TEST_CHECK_RETRY_DURATION));

        assertThat(
                        mHelper.handleGetDeviceCheckInStatusResponse(
                                response, mScheduler, /* fcmRegistrationToken= */ null))
                .isTrue();

        List<WorkInfo> actualWorks =
                Futures.getUnchecked(mWorkManager.getWorkInfosForUniqueWork(FCM_TOKEN_WORK_NAME));
        assertThat(actualWorks.size()).isEqualTo(1);
        WorkInfo actualWorkInfo = actualWorks.get(0);
        NetworkRequest networkRequest = actualWorkInfo.getConstraints().getRequiredNetworkRequest();
        assertNetworkRequestCapabilities(networkRequest);
        assertThat(actualWorkInfo.getInitialDelayMillis())
                .isEqualTo(FCM_TOKEN_WORKER_INITIAL_DELAY.toMillis());
    }

    @Test
    public void handleGetDeviceCheckInStatusResponse_retryCheckIn_durationIsNegative_shouldRetry()
            throws Exception {
        final GetDeviceCheckInStatusGrpcResponse response =
                createRetryResponse(
                        SystemClock.currentNetworkTimeClock()
                                .instant()
                                .plus(TEST_NEGATIVE_CHECK_RETRY_DURATION));

        assertThat(
                        mHelper.handleGetDeviceCheckInStatusResponse(
                                response,
                                mScheduler,
                                mTestApplication.getFcmRegistrationToken().get()))
                .isTrue();

        verify(mScheduler).scheduleRetryCheckInWork(eq(Duration.ZERO));
    }

    @Test
    public void getDeviceLockApexVersion_missingPackageName_shouldReturnZero() {
        assertThat(mHelper.getDeviceLockApexVersion("non_existent_package")).isEqualTo(0);
    }

    @Test
    public void getDeviceLockApexVersion_validPackageName_returnsVersion() {
        PackageInfo packageInfo = new PackageInfo();
        packageInfo.setLongVersionCode(APEX_VERSION);
        packageInfo.packageName = FAKE_APEX_PACKAGE;
        packageInfo.isApex = true;
        mShadowPackageManager.installPackage(packageInfo);
        assertThat(mHelper.getDeviceLockApexVersion(FAKE_APEX_PACKAGE)).isEqualTo(APEX_VERSION);
    }

    @Test
    public void getDeviceLocale_returnsDefaultLocale() {
        assertThat(mHelper.getDeviceLocale()).isEqualTo(DEVICE_LOCALE);
    }

    @Test
    public void getCheckInRequiredPackageState_enabledPackage_returnsEnabled() {
        installSystemPhonePackage(PackageManager.COMPONENT_ENABLED_STATE_ENABLED);

        assertThat(mHelper.getCheckInRequiredPackageState(PHONE_SYSTEM_APP_PACKAGE_NAME))
                .isEqualTo(AbstractDeviceCheckInHelper.CheckInRequiredPackageState.ENABLED);
    }

    @Test
    public void getCheckInRequiredPackageState_disablesPackage_returnsDisabled() {
        installSystemPhonePackage(PackageManager.COMPONENT_ENABLED_STATE_DISABLED_USER);

        assertThat(mHelper.getCheckInRequiredPackageState(PHONE_SYSTEM_APP_PACKAGE_NAME))
                .isEqualTo(AbstractDeviceCheckInHelper.CheckInRequiredPackageState.DISABLED);
    }

    @Test
    public void getCheckInRequiredPackageState_missingPackage_returnsUninstalled() {
        assertThat(mHelper.getCheckInRequiredPackageState(PHONE_SYSTEM_APP_PACKAGE_NAME))
                .isEqualTo(AbstractDeviceCheckInHelper.CheckInRequiredPackageState.UNINSTALLED);
    }

    @Test
    public void enableCheckInRequiredPackage_disabledPackage_enablesPackage() {
        installSystemPhonePackage(PackageManager.COMPONENT_ENABLED_STATE_DISABLED_USER);

        mHelper.enableCheckInRequiredPackage(
                PHONE_SYSTEM_APP_PACKAGE_NAME,
                AbstractDeviceCheckInHelper.CheckInRequiredPackageState.DISABLED);

        assertThat(mPackageManager.getApplicationEnabledSetting(PHONE_SYSTEM_APP_PACKAGE_NAME))
                .isEqualTo(COMPONENT_ENABLED_STATE_DEFAULT);
        verify(mMockDevicePolicyController)
                .disableUserControlForCheckInRequiredPackage(PHONE_SYSTEM_APP_PACKAGE_NAME);
    }

    @Test
    public void enableCheckInRequiredPackage_uninstalledPackage_reInstallsPackage() {
        mHelper.enableCheckInRequiredPackage(
                PHONE_SYSTEM_APP_PACKAGE_NAME,
                AbstractDeviceCheckInHelper.CheckInRequiredPackageState.DISABLED);

        // Actual behaviour for installing existing packages added in instrumentation tests
        verify(mMockDevicePolicyController)
                .disableUserControlForCheckInRequiredPackage(PHONE_SYSTEM_APP_PACKAGE_NAME);
    }

    @Test
    public void hasTelephonyFeature_featurePresent_returnsTrue() {
        mShadowPackageManager.setSystemFeature(
                PackageManager.FEATURE_TELEPHONY, /* supported= */ true);

        assertThat(mHelper.hasTelephonyFeature()).isTrue();
    }

    @Test
    public void hasTelephonyFeature_featureNotPresent_returnsFalse() {
        assertThat(mHelper.hasTelephonyFeature()).isFalse();
    }

    private void installSystemPhonePackage(int enabledSetting) {
        PackageInfo packageInfo = new PackageInfo();
        packageInfo.packageName = PHONE_SYSTEM_APP_PACKAGE_NAME;

        mShadowPackageManager.installPackage(packageInfo);
        mPackageManager.setApplicationEnabledSetting(packageInfo.packageName, enabledSetting, 0);
    }

    private void assertNetworkRequestCapabilities(NetworkRequest networkRequest) {
        assertThat(networkRequest.hasCapability(NET_CAPABILITY_NOT_RESTRICTED)).isTrue();
        assertThat(networkRequest.hasCapability(NET_CAPABILITY_TRUSTED)).isTrue();
        assertThat(networkRequest.hasCapability(NET_CAPABILITY_INTERNET)).isTrue();
        assertThat(networkRequest.hasCapability(NET_CAPABILITY_NOT_VPN)).isTrue();
    }

    private static GetDeviceCheckInStatusGrpcResponse createStopResponse() {
        return createMockResponse(STOP_CHECK_IN, /* nextCheckInTime= */ null, /* config= */ null);
    }

    private static GetDeviceCheckInStatusGrpcResponse createRetryResponse(Instant nextCheckInTime) {
        return createMockResponse(RETRY_CHECK_IN, nextCheckInTime, /* config= */ null);
    }

    private static GetDeviceCheckInStatusGrpcResponse createReadyResponse(
            ProvisioningConfiguration configuration) {
        return createMockResponse(READY_FOR_PROVISION, /* nextCheckInTime= */ null, configuration);
    }

    private static GetDeviceCheckInStatusGrpcResponse createMockResponse(
            @DeviceCheckInStatus int checkInStatus,
            @Nullable Instant nextCheckInTime,
            @Nullable ProvisioningConfiguration config) {
        GetDeviceCheckInStatusGrpcResponse response =
                Mockito.mock(GetDeviceCheckInStatusGrpcResponse.class);
        when(response.getDeviceCheckInStatus()).thenReturn(checkInStatus);
        when(response.isProvisioningMandatory()).thenReturn(IS_PROVISIONING_MANDATORY);
        if (nextCheckInTime != null) {
            when(response.getNextCheckInTime()).thenReturn(nextCheckInTime);
        }
        if (config != null) {
            when(response.getProvisioningConfig()).thenReturn(config);
        }
        return response;
    }
}