xref: /aosp_15_r20/cts/tests/tests/virtualdevice/common/src/android/virtualdevice/cts/common/VirtualDeviceRule.java (revision b7c941bb3fa97aba169d73cee0bed2de8ac964bf)

/*
 * 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 android.virtualdevice.cts.common;

import static android.content.pm.PackageManager.FEATURE_ACTIVITIES_ON_SECONDARY_DISPLAYS;
import static android.content.pm.PackageManager.FEATURE_FREEFORM_WINDOW_MANAGEMENT;

import static androidx.test.platform.app.InstrumentationRegistry.getInstrumentation;

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

import static org.junit.Assume.assumeFalse;
import static org.junit.Assume.assumeNotNull;
import static org.junit.Assume.assumeTrue;

import android.annotation.TargetApi;
import android.app.Activity;
import android.app.ActivityOptions;
import android.app.UiAutomation;
import android.companion.AssociationInfo;
import android.companion.AssociationRequest;
import android.companion.virtual.VirtualDeviceManager;
import android.companion.virtual.VirtualDeviceManager.VirtualDevice;
import android.companion.virtual.VirtualDeviceParams;
import android.content.ComponentName;
import android.content.Context;
import android.content.Intent;
import android.hardware.display.DisplayManager;
import android.hardware.display.VirtualDisplay;
import android.hardware.display.VirtualDisplayConfig;
import android.media.ImageReader;
import android.os.Bundle;
import android.platform.test.flag.junit.DeviceFlagsValueProvider;
import android.server.wm.Condition;
import android.server.wm.WindowManagerState;
import android.server.wm.WindowManagerStateHelper;
import android.view.Display;
import android.view.Surface;

import androidx.annotation.NonNull;
import androidx.annotation.Nullable;
import androidx.core.os.BuildCompat;

import com.android.compatibility.common.util.AdoptShellPermissionsRule;
import com.android.compatibility.common.util.FeatureUtil;

import org.junit.rules.ExternalResource;
import org.junit.rules.RuleChain;
import org.junit.rules.TestRule;
import org.junit.runner.Description;
import org.junit.runners.model.Statement;

import java.util.ArrayList;
import java.util.Set;
import java.util.function.BooleanSupplier;
import java.util.function.Supplier;

/**
 * A test rule that allows for testing VDM and virtual device features.
 */
@TargetApi(34)
public class VirtualDeviceRule implements TestRule {

    public static final VirtualDeviceParams DEFAULT_VIRTUAL_DEVICE_PARAMS =
            new VirtualDeviceParams.Builder().build();

    public static final String DEFAULT_VIRTUAL_DISPLAY_NAME = "testVirtualDisplay";
    public static final int DEFAULT_VIRTUAL_DISPLAY_WIDTH = 640;
    public static final int DEFAULT_VIRTUAL_DISPLAY_HEIGHT = 480;
    public static final int DEFAULT_VIRTUAL_DISPLAY_DPI = 240;

    public static final ComponentName BLOCKED_ACTIVITY_COMPONENT =
            new ComponentName("android", "com.android.internal.app.BlockedAppStreamingActivity");

    private RuleChain mRuleChain;
    private final FakeAssociationRule mFakeAssociationRule;
    private final VirtualDeviceTrackerRule mTrackerRule = new VirtualDeviceTrackerRule();

    private final Context mContext = getInstrumentation().getTargetContext();
    private final VirtualDeviceManager mVirtualDeviceManager =
            mContext.getSystemService(VirtualDeviceManager.class);
    private final WindowManagerStateHelper mWmState = new WindowManagerStateHelper();

    /** A default virtual device for tests that only use the rule to access VDM functionality. */
    private VirtualDevice mDefaultVirtualDevice = null;

    /** Creates a rule with the required permissions for creating virtual devices and displays. */
    public static VirtualDeviceRule createDefault() {
        return new VirtualDeviceRule(AssociationRequest.DEVICE_PROFILE_APP_STREAMING);
    }

    /** Creates a rule with an explicit device profile. */
    public static VirtualDeviceRule withDeviceProfile(String deviceProfile) {
        return new VirtualDeviceRule(deviceProfile);
    }

    /** Creates a rule with any additional permission needed for the specific test. */
    public static VirtualDeviceRule withAdditionalPermissions(String... additionalPermissions) {
        return new VirtualDeviceRule(AssociationRequest.DEVICE_PROFILE_APP_STREAMING,
                additionalPermissions);
    }

    private VirtualDeviceRule(String deviceProfile, String... permissions) {
        mFakeAssociationRule = new FakeAssociationRule(deviceProfile);
        mRuleChain = RuleChain
                .outerRule(mFakeAssociationRule)
                .around(DeviceFlagsValueProvider.createCheckFlagsRule())
                .around(new AdoptShellPermissionsRule(
                        getInstrumentation().getUiAutomation(), permissions))
                .around(mTrackerRule);
    }

    /** Creates a rule with virtual camera support check before test execution. */
    public VirtualDeviceRule withVirtualCameraSupportCheck() {
        mRuleChain = mRuleChain.around(new VirtualCameraSupportRule(this));
        return this;
    }

    @Override
    public Statement apply(final Statement base, final Description description) {
        assumeNotNull(mVirtualDeviceManager);
        return mRuleChain.apply(base, description);
    }

    /**
     * Returns a default virtual device.
     */
    public VirtualDevice getDefaultVirtualDevice() {
        if (mDefaultVirtualDevice == null) {
            mDefaultVirtualDevice = createManagedVirtualDevice();
        }
        return mDefaultVirtualDevice;
    }

    /**
     * Returns the VirtualDevice object for the given deviceId
     */
    public android.companion.virtual.VirtualDevice getVirtualDevice(int deviceId) {
        if (BuildCompat.isAtLeastV()) {
            return mVirtualDeviceManager.getVirtualDevice(deviceId);
        } else {
            return mVirtualDeviceManager.getVirtualDevices().stream()
                    .filter(device -> device.getDeviceId() == deviceId).findFirst().orElse(null);
        }
    }

    /**
     * Creates a virtual device with default params that will be automatically closed when the
     * test is torn down.
     */
    @NonNull
    public VirtualDevice createManagedVirtualDevice() {
        return createManagedVirtualDevice(DEFAULT_VIRTUAL_DEVICE_PARAMS);
    }

    /**
     * Creates a virtual device with the given params that will be automatically closed when the
     * test is torn down.
     */
    @NonNull
    public VirtualDevice createManagedVirtualDevice(@NonNull VirtualDeviceParams params) {
        final VirtualDevice virtualDevice = mVirtualDeviceManager.createVirtualDevice(
                mFakeAssociationRule.getAssociationInfo().getId(), params);
        mTrackerRule.mVirtualDevices.add(virtualDevice);
        return virtualDevice;
    }

    /**
     * Creates a virtual display associated with the given device that will be automatically
     * released when the test is torn down.
     */
    @Nullable
    public VirtualDisplay createManagedVirtualDisplay(@NonNull VirtualDevice virtualDevice) {
        return createManagedVirtualDisplay(virtualDevice,
                createDefaultVirtualDisplayConfigBuilder());
    }

    /**
     * Creates a virtual display associated with the given device and flags that will be
     * automatically released when the test is torn down.
     */
    @Nullable
    public VirtualDisplay createManagedVirtualDisplayWithFlags(
            @NonNull VirtualDevice virtualDevice, int flags) {
        return createManagedVirtualDisplay(virtualDevice,
                createDefaultVirtualDisplayConfigBuilder().setFlags(flags));
    }

    /**
     * Creates a virtual display associated with the given device and config that will be
     * automatically released when the test is torn down.
     */
    @Nullable
    public VirtualDisplay createManagedVirtualDisplay(@NonNull VirtualDevice virtualDevice,
            @NonNull VirtualDisplayConfig.Builder builder) {
        return createManagedVirtualDisplay(virtualDevice, builder, /* callback= */ null);
    }

    /**
     * Creates a virtual display associated with the given device and config that will be
     * automatically released when the test is torn down.
     */
    @Nullable
    public VirtualDisplay createManagedVirtualDisplay(@NonNull VirtualDevice virtualDevice,
            @NonNull VirtualDisplayConfig.Builder builder,
            @Nullable VirtualDisplay.Callback callback) {
        VirtualDisplayConfig config = builder.build();
        final Surface surface = prepareSurface(config.getWidth(), config.getHeight());
        final VirtualDisplay virtualDisplay = virtualDevice.createVirtualDisplay(
                builder.setSurface(surface).build(), Runnable::run, callback);
        if (virtualDisplay != null) {
            assertDisplayExists(virtualDisplay.getDisplay().getDisplayId());
            // There's no need to track managed virtual displays to have them released on tear-down
            // because they will be released automatically when the VirtualDevice is closed.
        }
        return virtualDisplay;
    }

    /**
     * Creates a virtual display not associated with the any virtual device that will be
     * automatically released when the test is torn down.
     */
    @Nullable
    public VirtualDisplay createManagedUnownedVirtualDisplay() {
        return createManagedUnownedVirtualDisplay(createTrustedVirtualDisplayConfigBuilder());
    }

    /**
     * Creates a virtual display not associated with the any virtual device with the given flags
     * that will be automatically released when the test is torn down.
     */
    @Nullable
    public VirtualDisplay createManagedUnownedVirtualDisplayWithFlags(int flags) {
        return createManagedUnownedVirtualDisplay(
                createDefaultVirtualDisplayConfigBuilder().setFlags(flags));
    }

    /**
     * Creates a virtual display not associated with the any virtual device with the given config
     * that will be automatically released when the test is torn down.
     */
    @Nullable
    public VirtualDisplay createManagedUnownedVirtualDisplay(
            @NonNull VirtualDisplayConfig.Builder builder) {
        VirtualDisplayConfig config = builder.build();
        final Surface surface = prepareSurface(config.getWidth(), config.getHeight());
        final VirtualDisplay virtualDisplay =
                mContext.getSystemService(DisplayManager.class).createVirtualDisplay(
                        builder.setSurface(surface).build());
        if (virtualDisplay != null) {
            assertDisplayExists(virtualDisplay.getDisplay().getDisplayId());
            mTrackerRule.mVirtualDisplays.add(virtualDisplay);
        }
        return virtualDisplay;
    }

    /**
     * Default config for virtual display creation, with a predefined name, dimensions and an empty
     * surface.
     */
    @NonNull
    public static VirtualDisplayConfig.Builder createDefaultVirtualDisplayConfigBuilder() {
        return createDefaultVirtualDisplayConfigBuilder(
                DEFAULT_VIRTUAL_DISPLAY_WIDTH, DEFAULT_VIRTUAL_DISPLAY_HEIGHT);
    }

    /**
     * Config for trusted virtual display creation, which applies the flags
     * {@link DisplayManager#VIRTUAL_DISPLAY_FLAG_TRUSTED},
     * {@link DisplayManager#VIRTUAL_DISPLAY_FLAG_PUBLIC} and
     * {@link DisplayManager#VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY}. A predefined name and
     * dimensions and an empty surface are used.
     */
    @NonNull
    public static VirtualDisplayConfig.Builder createTrustedVirtualDisplayConfigBuilder() {
        return createDefaultVirtualDisplayConfigBuilder().setFlags(
                DisplayManager.VIRTUAL_DISPLAY_FLAG_PUBLIC
                | DisplayManager.VIRTUAL_DISPLAY_FLAG_TRUSTED
                | DisplayManager.VIRTUAL_DISPLAY_FLAG_OWN_CONTENT_ONLY);
    }

    /**
     * Default config for virtual display creation with custom dimensions.
     */
    @NonNull
    public static VirtualDisplayConfig.Builder createDefaultVirtualDisplayConfigBuilder(
            int width, int height) {
        return new VirtualDisplayConfig.Builder(
                DEFAULT_VIRTUAL_DISPLAY_NAME, width, height, DEFAULT_VIRTUAL_DISPLAY_DPI);
    }

    /**
     * Blocks until the display with the given ID is available.
     */
    public void assertDisplayExists(int displayId) {
        waitAndAssertWindowManagerState("Waiting for display to be available",
                () -> mWmState.getDisplay(displayId) != null);
    }

    /**
     * Blocks until the display with the given ID is removed.
     */
    public void assertDisplayDoesNotExist(int displayId) {
        waitAndAssertWindowManagerState("Waiting for display to be removed",
                () -> mWmState.getDisplay(displayId) == null);
    }

    /** Returns the WM state helper. */
    public WindowManagerStateHelper getWmState() {
        return mWmState;
    }

    /** Creates a new CDM association. */
    public AssociationInfo createManagedAssociation() {
        return mFakeAssociationRule.createManagedAssociation();
    }

    /** Drops the current CDM association. */
    public void dropCompanionDeviceAssociation() {
        mFakeAssociationRule.disassociate();
    }

    /**
     * Temporarily assumes the given permissions and executes the given supplier. Reverts any
     * permissions currently held after the execution.
     */
    public <T> T runWithTemporaryPermission(Supplier<T> supplier, String... permissions) {
        UiAutomation uiAutomation = getInstrumentation().getUiAutomation();
        final Set<String> currentPermissions = uiAutomation.getAdoptedShellPermissions();
        uiAutomation.adoptShellPermissionIdentity(permissions);
        try {
            return supplier.get();
        } finally {
            // Revert the permissions needed for the test again.
            uiAutomation.adoptShellPermissionIdentity(
                    currentPermissions.toArray(new String[0]));
        }
    }

    /**
     * Temporarily drops any permissions and executes the given supplier. Reverts any permissions
     * currently held after the execution.
     */
    public <T> T runWithoutPermissions(Supplier<T> supplier) {
        UiAutomation uiAutomation = getInstrumentation().getUiAutomation();
        final Set<String> currentPermissions = uiAutomation.getAdoptedShellPermissions();
        uiAutomation.dropShellPermissionIdentity();
        try {
            return supplier.get();
        } finally {
            // Revert the permissions needed for the test again.
            uiAutomation.adoptShellPermissionIdentity(
                    currentPermissions.toArray(new String[0]));
        }
    }

    /**
     * Sends the given intent to the given virtual display.
     */
    public void sendIntentToDisplay(Intent intent, VirtualDisplay virtualDisplay) {
        sendIntentToDisplay(intent, virtualDisplay.getDisplay().getDisplayId());
    }

    /**
     * Sends the given intent to the given display.
     */
    public void sendIntentToDisplay(Intent intent, int displayId) {
        assumeActivityLaunchSupported(displayId);
        mContext.startActivity(intent, createActivityOptions(displayId));
    }

    /**
     * Starts the activity for the given class on the given virtual display and blocks until it is
     * successfully launched there. The activity will be finished after the test run.
     */
    public <T extends Activity> T startActivityOnDisplaySync(
            VirtualDisplay virtualDisplay, Class<T> clazz) {
        final int displayId = virtualDisplay.getDisplay().getDisplayId();
        return startActivityOnDisplaySync(displayId, clazz);
    }

    /**
     * Starts the activity for the given class on the given display and blocks until it is
     * successfully launched there. The activity will be finished after the test run.
     */
    public <T extends Activity> T startActivityOnDisplaySync(int displayId, Class<T> clazz) {
        return startActivityOnDisplaySync(displayId, new Intent(mContext, clazz)
                .addFlags(Intent.FLAG_ACTIVITY_NEW_TASK | Intent.FLAG_ACTIVITY_CLEAR_TASK));
    }

    /**
     * Starts the activity for the given intent on the given virtual display and blocks until it is
     * successfully launched there. The activity will be finished after the test run.
     */
    public <T extends Activity> T startActivityOnDisplaySync(
            VirtualDisplay virtualDisplay, Intent intent) {
        return startActivityOnDisplaySync(virtualDisplay.getDisplay().getDisplayId(), intent);
    }

    /**
     * Starts the activity for the given intent on the given display and blocks until it is
     * successfully launched there. The activity will be finished after the test run.
     */
    public <T extends Activity> T startActivityOnDisplaySync(int displayId, Intent intent) {
        assumeActivityLaunchSupported(displayId);
        T activity = (T) getInstrumentation()
                .startActivitySync(intent, createActivityOptions(displayId));
        mTrackerRule.mActivities.add(activity);
        return activity;
    }

    /**
     * Creates activity options for launching activities on the given virtual display.
     */
    public static Bundle createActivityOptions(VirtualDisplay virtualDisplay) {
        return createActivityOptions(virtualDisplay.getDisplay().getDisplayId());
    }

    /**
     * Creates activity options for launching activities on the given display.
     */
    public static Bundle createActivityOptions(int displayId) {
        return ActivityOptions.makeBasic().setLaunchDisplayId(displayId).toBundle();
    }

    /**
     * Skips the test if the device doesn't support virtual displays that can host activities.
     */
    public void assumeActivityLaunchSupported(int displayId) {
        if (displayId != Display.DEFAULT_DISPLAY) {
            assumeTrue(FeatureUtil.hasSystemFeature(FEATURE_ACTIVITIES_ON_SECONDARY_DISPLAYS));
            // TODO(b/261155110): Re-enable once freeform mode is supported on virtual displays.
            assumeFalse(FeatureUtil.hasSystemFeature(FEATURE_FREEFORM_WINDOW_MANAGEMENT));
        }
    }

    /**
     * Blocks until the given activity is in resumed state.
     */
    public void waitAndAssertActivityResumed(ComponentName componentName) {
        waitAndAssertWindowManagerState("Waiting for activity to be resumed",
                () -> mWmState.hasActivityState(componentName, WindowManagerState.STATE_RESUMED));
    }

    /**
     * Blocks until the given activity is in resumed state on the given display.
     */
    public void waitAndAssertActivityResumed(ComponentName componentName, int displayId) {
        waitAndAssertActivityResumed(componentName);
        mWmState.assertResumedActivities("Activity must be on display " + displayId,
                mapping -> mapping.put(displayId, componentName));
    }

    /**
     * Blocks until the given activity is gone.
     */
    public void waitAndAssertActivityRemoved(ComponentName componentName) {
        waitAndAssertWindowManagerState("Waiting for activity to be removed",
                () -> !mWmState.containsActivity(componentName));
    }

    /**
     * Override the default retry limit of WindowManagerStateHelper.
     * Destroying activities on virtual displays and destroying the virtual displays themselves
     * takes longer than the default timeout of 5s.
     */
    private void waitAndAssertWindowManagerState(
            String message, BooleanSupplier waitCondition) {
        final Condition<String> condition =
                new Condition<>(message, () -> {
                    mWmState.computeState();
                    return waitCondition.getAsBoolean();
                });
        condition.setRetryLimit(10);
        assertThat(Condition.waitFor(condition)).isTrue();
    }

    private Surface prepareSurface(int width, int height) {
        ImageReader imageReader = new ImageReader.Builder(width, height).build();
        mTrackerRule.mImageReaders.add(imageReader);
        return imageReader.getSurface();
    }

    /**
     * Internal rule that tracks all created virtual devices and displays and ensures they are
     * properly closed and released after the test.
     */
    private static final class VirtualDeviceTrackerRule extends ExternalResource {

        final ArrayList<VirtualDevice> mVirtualDevices = new ArrayList<>();
        final ArrayList<VirtualDisplay> mVirtualDisplays = new ArrayList<>();
        final ArrayList<Activity> mActivities = new ArrayList<>();
        final ArrayList<ImageReader> mImageReaders = new ArrayList<>();

        @Override
        protected void after() {
            for (Activity activity : mActivities) {
                activity.finish();
            }
            mActivities.clear();

            for (VirtualDevice virtualDevice : mVirtualDevices) {
                virtualDevice.close();
            }
            mVirtualDevices.clear();
            for (VirtualDisplay virtualDisplay : mVirtualDisplays) {
                virtualDisplay.release();
            }
            mVirtualDisplays.clear();
            for (ImageReader imageReader : mImageReaders) {
                imageReader.close();
            }
            mImageReaders.clear();
            super.after();
        }
    }

}