tests/servicetests/src/com/android/ondevicepersonalization/services/federatedcompute/OdpFederatedComputeJobServiceTest.java

/*
 * Copyright (C) 2023 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.ondevicepersonalization.services.federatedcompute;

import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyBoolean;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;

import android.app.job.JobParameters;
import android.content.Context;
import android.federatedcompute.FederatedComputeManager;

import androidx.test.core.app.ApplicationProvider;

import com.android.dx.mockito.inline.extended.ExtendedMockito;
import com.android.ondevicepersonalization.services.OnDevicePersonalizationExecutors;
import com.android.ondevicepersonalization.services.PhFlagsTestUtil;

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

import org.junit.Before;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;
import org.mockito.MockitoSession;
import org.mockito.quality.Strictness;

@RunWith(JUnit4.class)
public class OdpFederatedComputeJobServiceTest {
    private final Context mContext = ApplicationProvider.getApplicationContext();
    private OdpFederatedComputeJobService mSpyService;

    @Before
    public void setup() throws Exception {
        PhFlagsTestUtil.setUpDeviceConfigPermissions();
        PhFlagsTestUtil.disableGlobalKillSwitch();
        mSpyService = spy(new OdpFederatedComputeJobService());
    }

    @Test
    public void onStartJobTest() {
        MockitoSession session =
                ExtendedMockito.mockitoSession()
                        .spyStatic(OnDevicePersonalizationExecutors.class)
                        .strictness(Strictness.LENIENT)
                        .startMocking();
        try {
            FederatedComputeManager mockManager = mock(FederatedComputeManager.class);
            doNothing().when(mSpyService).jobFinished(any(), anyBoolean());
            doNothing().when(mockManager).schedule(any(), any(), any());
            doReturn(mockManager).when(mSpyService).getSystemService(FederatedComputeManager.class);
            ExtendedMockito.doReturn(MoreExecutors.newDirectExecutorService())
                    .when(OnDevicePersonalizationExecutors::getBackgroundExecutor);
            ExtendedMockito.doReturn(MoreExecutors.newDirectExecutorService())
                    .when(OnDevicePersonalizationExecutors::getLightweightExecutor);

            boolean result = mSpyService.onStartJob(mock(JobParameters.class));
            assertTrue(result);
            verify(mSpyService, times(1)).jobFinished(any(), eq(false));
            verify(mockManager, times(1)).schedule(any(), any(), any());
        } finally {
            session.finishMocking();
        }
    }

    @Test
    public void onStartJobNullFederatedComputeManagerTest() {
        MockitoSession session = ExtendedMockito.mockitoSession().spyStatic(
                OnDevicePersonalizationExecutors.class).strictness(
                Strictness.LENIENT).startMocking();
        try {
            doReturn(null).when(mSpyService).getSystemService(FederatedComputeManager.class);
            ExtendedMockito.doReturn(MoreExecutors.newDirectExecutorService()).when(
                    OnDevicePersonalizationExecutors::getBackgroundExecutor);
            ExtendedMockito.doReturn(MoreExecutors.newDirectExecutorService()).when(
                    OnDevicePersonalizationExecutors::getLightweightExecutor);

            boolean result = mSpyService.onStartJob(mock(JobParameters.class));
            assertTrue(result);
        } catch (Exception e) {
            fail("An unexpected exception was thrown: " + e.getMessage());
        } finally {
            session.finishMocking();
        }
    }

    @Test
    public void onStartJobTestKillSwitchEnabled() {
        PhFlagsTestUtil.enableGlobalKillSwitch();
        MockitoSession session =
                ExtendedMockito.mockitoSession().strictness(Strictness.LENIENT).startMocking();
        try {
            FederatedComputeManager mockManager = mock(FederatedComputeManager.class);
            doNothing().when(mSpyService).jobFinished(any(), anyBoolean());
            boolean result = mSpyService.onStartJob(mock(JobParameters.class));
            assertTrue(result);
            verify(mSpyService, times(1)).jobFinished(any(), eq(false));
            verify(mockManager, times(0)).schedule(any(), any(), any());
        } finally {
            session.finishMocking();
        }
    }

    @Test
    public void onStopJobTest() {
        MockitoSession session =
                ExtendedMockito.mockitoSession().strictness(Strictness.LENIENT).startMocking();
        try {
            assertTrue(mSpyService.onStopJob(mock(JobParameters.class)));
        } finally {
            session.finishMocking();
        }
    }
}