From 5789bf94dd6506d430e0487e50c2d7602c71f6bf Mon Sep 17 00:00:00 2001
From: Jan Sebechlebsky <jsebechlebsky@google.com>
Date: Tue, 13 Dec 2022 13:55:01 +0100
Subject: [PATCH] Add context-aware AudioTrack construction APIs.

This cl also introduces new device policy type for audio.
The POLICY_TYPE_AUDIO is used to specialize AudioManager
behavior for device-specific context - if the context
belongs to VirtualDevice instance with custom device policy
for audio, it will try to extract the device specific session
id for that device from VDM. The device specific session aid
is used to aid virtual device audio routing.

The context-aware constructor and builder will be extended
in follow-up cls to extract AttributionSource from the Context
and pass it to the native layer.

The policy will be used to specialize AudioManager behavior for
device-specific context.

Bug: 261698699
Bug: 249777386
Test: atest VirtualDeviceParamsTest AudioTrackUnitTest AudioTrackTest
Change-Id: I558b508f06e8098a8fbebaa0560557af06cdd0f7
---
 core/api/current.txt                          |   1 +
 core/api/system-current.txt                   |   1 +
 .../virtual/VirtualDeviceManager.java         |  38 +++++
 .../virtual/VirtualDeviceParams.java          |  17 +-
 media/java/android/media/AudioTrack.java      |  79 ++++++++--
 media/tests/MediaFrameworkTest/Android.bp     |   9 +-
 .../MediaFrameworkTest/AndroidManifest.xml    |   3 +-
 .../unit/AudioTrackUnitTest.java              | 145 ++++++++++++++++++
 .../virtual/VirtualDeviceParamsTest.java      |   3 +
 9 files changed, 281 insertions(+), 15 deletions(-)
 create mode 100644 media/tests/MediaFrameworkTest/src/com/android/mediaframeworktest/unit/AudioTrackUnitTest.java

diff --git a/core/api/current.txt b/core/api/current.txt
index 39ac0d2ebad2c..ad2239bfa0579 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -21182,6 +21182,7 @@ package android.media {
     method @NonNull public android.media.AudioTrack.Builder setAudioAttributes(@NonNull android.media.AudioAttributes) throws java.lang.IllegalArgumentException;
     method @NonNull public android.media.AudioTrack.Builder setAudioFormat(@NonNull android.media.AudioFormat) throws java.lang.IllegalArgumentException;
     method @NonNull public android.media.AudioTrack.Builder setBufferSizeInBytes(@IntRange(from=0) int) throws java.lang.IllegalArgumentException;
+    method @NonNull public android.media.AudioTrack.Builder setContext(@NonNull android.content.Context);
     method @NonNull public android.media.AudioTrack.Builder setEncapsulationMode(int);
     method @NonNull public android.media.AudioTrack.Builder setOffloadedPlayback(boolean);
     method @NonNull public android.media.AudioTrack.Builder setPerformanceMode(int);
diff --git a/core/api/system-current.txt b/core/api/system-current.txt
index fb742607d9d25..e6e140cf1ae46 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3035,6 +3035,7 @@ package android.companion.virtual {
     field public static final int LOCK_STATE_DEFAULT = 0; // 0x0
     field public static final int NAVIGATION_POLICY_DEFAULT_ALLOWED = 0; // 0x0
     field public static final int NAVIGATION_POLICY_DEFAULT_BLOCKED = 1; // 0x1
+    field public static final int POLICY_TYPE_AUDIO = 1; // 0x1
     field public static final int POLICY_TYPE_SENSORS = 0; // 0x0
     field public static final int RECENTS_POLICY_ALLOW_IN_HOST_DEVICE_RECENTS = 1; // 0x1
   }
diff --git a/core/java/android/companion/virtual/VirtualDeviceManager.java b/core/java/android/companion/virtual/VirtualDeviceManager.java
index dba7c8e630c80..b0292613207f6 100644
--- a/core/java/android/companion/virtual/VirtualDeviceManager.java
+++ b/core/java/android/companion/virtual/VirtualDeviceManager.java
@@ -16,6 +16,8 @@
 
 package android.companion.virtual;
 
+import static android.media.AudioManager.AUDIO_SESSION_ID_GENERATE;
+
 import android.annotation.CallbackExecutor;
 import android.annotation.IntDef;
 import android.annotation.IntRange;
@@ -252,6 +254,42 @@ public final class VirtualDeviceManager {
         }
     }
 
+    /**
+     * Returns device-specific audio session id for audio playback.
+     *
+     * @param deviceId - id of the virtual audio device
+     * @return Device specific session id to be used for audio playback (see
+     *     {@link android.media.AudioManager.generateAudioSessionId}) if virtual device has
+     *     {@link VirtualDeviceParams.POLICY_TYPE_AUDIO} set to
+     *     {@link VirtualDeviceParams.DEVICE_POLICY_CUSTOM} and Virtual Audio Device
+     *     is configured in context-aware mode.
+     *     Otherwise {@link AUDIO_SESSION_ID_GENERATE} constant is returned.
+     * @hide
+     */
+    public int getAudioPlaybackSessionId(int deviceId) {
+        //TODO - Return session id rerouted to VirtualAudioDevice if the VirtualAudioDevice
+        //is configured to operate in context-aware mode.
+        return AUDIO_SESSION_ID_GENERATE;
+    }
+
+    /**
+     * Returns device-specific audio session id for audio recording.
+     *
+     * @param deviceId - id of the virtual audio device
+     * @return Device specific session id to be used for audio recording (see
+     *     {@link android.media.AudioManager.generateAudioSessionId}) if virtual device has
+     *     {@link VirtualDeviceParams.POLICY_TYPE_AUDIO} set to
+     *     {@link VirtualDeviceParams.DEVICE_POLICY_CUSTOM} and Virtual Audio Device
+     *     is configured in context-aware mode.
+     *     Otherwise {@link AUDIO_SESSION_ID_GENERATE} constant is returned.
+     * @hide
+     */
+    public int getAudioRecordingSessionId(int deviceId) {
+        //TODO - Return session id corresponding to VirtualAudioDevice injection if the
+        // VirtualAudioDevice is configured to operate in context-aware mode.
+        return AUDIO_SESSION_ID_GENERATE;
+    }
+
     /**
      * A virtual device has its own virtual display, audio output, microphone, and camera etc. The
      * creator of a virtual device can take the output from the virtual display and stream it over
diff --git a/core/java/android/companion/virtual/VirtualDeviceParams.java b/core/java/android/companion/virtual/VirtualDeviceParams.java
index be77f2b7e0f5e..f59a7a9b66683 100644
--- a/core/java/android/companion/virtual/VirtualDeviceParams.java
+++ b/core/java/android/companion/virtual/VirtualDeviceParams.java
@@ -129,7 +129,7 @@ public final class VirtualDeviceParams implements Parcelable {
      * a given policy type.
      * @hide
      */
-    @IntDef(prefix = "POLICY_TYPE_",  value = {POLICY_TYPE_SENSORS})
+    @IntDef(prefix = "POLICY_TYPE_",  value = {POLICY_TYPE_SENSORS, POLICY_TYPE_AUDIO})
     @Retention(RetentionPolicy.SOURCE)
     @Target({ElementType.TYPE_PARAMETER, ElementType.TYPE_USE})
     public @interface PolicyType {}
@@ -147,6 +147,21 @@ public final class VirtualDeviceParams implements Parcelable {
      */
     public static final int POLICY_TYPE_SENSORS = 0;
 
+    /**
+     * Tells the audio framework whether to configure the players ({@link android.media.AudioTrack},
+     * {@link android.media.MediaPlayer}, {@link android.media.SoundPool} and recorders
+     * {@link android.media.AudioRecord}) to use specific session ids re-routed to
+     * VirtualAudioDevice.
+     *
+     * <ul>
+     *     <li>{@link #DEVICE_POLICY_DEFAULT}: fall back to default session id handling.
+     *     <li>{@link #DEVICE_POLICY_CUSTOM}: audio framework will assign device specific session
+     *     ids to players and recorders constructed within device context. The session ids are
+     *     used to re-route corresponding audio streams to VirtualAudioDevice.
+     * <ul/>
+     */
+    public static final int POLICY_TYPE_AUDIO = 1;
+
     /** @hide */
     @IntDef(flag = true, prefix = "RECENTS_POLICY_",
             value = {RECENTS_POLICY_ALLOW_IN_HOST_DEVICE_RECENTS})
diff --git a/media/java/android/media/AudioTrack.java b/media/java/android/media/AudioTrack.java
index c2c752ed394e0..203429e1679ec 100644
--- a/media/java/android/media/AudioTrack.java
+++ b/media/java/android/media/AudioTrack.java
@@ -16,6 +16,11 @@
 
 package android.media;
 
+import static android.companion.virtual.VirtualDeviceManager.DEVICE_ID_DEFAULT;
+import static android.companion.virtual.VirtualDeviceParams.DEVICE_POLICY_DEFAULT;
+import static android.companion.virtual.VirtualDeviceParams.POLICY_TYPE_AUDIO;
+import static android.media.AudioManager.AUDIO_SESSION_ID_GENERATE;
+
 import android.annotation.CallbackExecutor;
 import android.annotation.FloatRange;
 import android.annotation.IntDef;
@@ -25,7 +30,9 @@ import android.annotation.Nullable;
 import android.annotation.RequiresPermission;
 import android.annotation.SystemApi;
 import android.annotation.TestApi;
+import android.companion.virtual.VirtualDeviceManager;
 import android.compat.annotation.UnsupportedAppUsage;
+import android.content.Context;
 import android.media.audiopolicy.AudioMix;
 import android.media.audiopolicy.AudioMixingRule;
 import android.media.audiopolicy.AudioPolicy;
@@ -545,7 +552,7 @@ public class AudioTrack extends PlayerBase
     /**
      * Audio session ID
      */
-    private int mSessionId = AudioManager.AUDIO_SESSION_ID_GENERATE;
+    private int mSessionId = AUDIO_SESSION_ID_GENERATE;
     /**
      * HW_AV_SYNC track AV Sync Header
      */
@@ -645,7 +652,7 @@ public class AudioTrack extends PlayerBase
             int bufferSizeInBytes, int mode)
     throws IllegalArgumentException {
         this(streamType, sampleRateInHz, channelConfig, audioFormat,
-                bufferSizeInBytes, mode, AudioManager.AUDIO_SESSION_ID_GENERATE);
+                bufferSizeInBytes, mode, AUDIO_SESSION_ID_GENERATE);
     }
 
     /**
@@ -749,12 +756,12 @@ public class AudioTrack extends PlayerBase
     public AudioTrack(AudioAttributes attributes, AudioFormat format, int bufferSizeInBytes,
             int mode, int sessionId)
                     throws IllegalArgumentException {
-        this(attributes, format, bufferSizeInBytes, mode, sessionId, false /*offload*/,
-                ENCAPSULATION_MODE_NONE, null /* tunerConfiguration */);
+        this(null /* context */, attributes, format, bufferSizeInBytes, mode, sessionId,
+                false /*offload*/, ENCAPSULATION_MODE_NONE, null /* tunerConfiguration */);
     }
 
-    private AudioTrack(AudioAttributes attributes, AudioFormat format, int bufferSizeInBytes,
-            int mode, int sessionId, boolean offload, int encapsulationMode,
+    private AudioTrack(@Nullable Context context, AudioAttributes attributes, AudioFormat format,
+            int bufferSizeInBytes, int mode, int sessionId, boolean offload, int encapsulationMode,
             @Nullable TunerConfiguration tunerConfiguration)
                     throws IllegalArgumentException {
         super(attributes, AudioPlaybackConfiguration.PLAYER_TYPE_JAM_AUDIOTRACK);
@@ -817,7 +824,13 @@ public class AudioTrack extends PlayerBase
 
         int[] sampleRate = new int[] {mSampleRate};
         int[] session = new int[1];
-        session[0] = sessionId;
+        if (sessionId == AUDIO_SESSION_ID_GENERATE) {
+            // If there's no specific session id requested, try to get one from context.
+            session[0] = getSessionIdForContext(context);
+        } else {
+            session[0] = sessionId;
+        }
+
         // native initialization
         int initResult = native_setup(new WeakReference<AudioTrack>(this), mAttributes,
                 sampleRate, mChannelMask, mChannelIndexMask, mAudioFormat,
@@ -1028,11 +1041,12 @@ public class AudioTrack extends PlayerBase
      * <br>Offload is false by default.
      */
     public static class Builder {
+        private Context mContext;
         private AudioAttributes mAttributes;
         private AudioFormat mFormat;
         private int mBufferSizeInBytes;
         private int mEncapsulationMode = ENCAPSULATION_MODE_NONE;
-        private int mSessionId = AudioManager.AUDIO_SESSION_ID_GENERATE;
+        private int mSessionId = AUDIO_SESSION_ID_GENERATE;
         private int mMode = MODE_STREAM;
         private int mPerformanceMode = PERFORMANCE_MODE_NONE;
         private boolean mOffload = false;
@@ -1045,6 +1059,19 @@ public class AudioTrack extends PlayerBase
         public Builder() {
         }
 
+        /**
+         * Sets the context the track belongs to. This context will be used to pull information,
+         * such as {@link android.content.AttributionSource} and device specific audio session ids,
+         * which will be associated with the {@link AudioTrack}. However, the context itself will
+         * not be retained by the {@link AudioTrack}.
+         * @param context a non-null {@link Context} instance
+         * @return the same Builder instance.
+         */
+        public @NonNull Builder setContext(@NonNull Context context) {
+            mContext = Objects.requireNonNull(context);
+            return this;
+        }
+
         /**
          * Sets the {@link AudioAttributes}.
          * @param attributes a non-null {@link AudioAttributes} instance that describes the audio
@@ -1152,6 +1179,10 @@ public class AudioTrack extends PlayerBase
 
         /**
          * Sets the session ID the {@link AudioTrack} will be attached to.
+         *
+         * Note, that if there's a device specific session id asociated with the context, explicitly
+         * setting a session id using this method will override it
+         * (see {@link Builder#setContext(Context)}).
          * @param sessionId a strictly positive ID number retrieved from another
          *     <code>AudioTrack</code> via {@link AudioTrack#getAudioSessionId()} or allocated by
          *     {@link AudioManager} via {@link AudioManager#generateAudioSessionId()}, or
@@ -1161,7 +1192,7 @@ public class AudioTrack extends PlayerBase
          */
         public @NonNull Builder setSessionId(@IntRange(from = 1) int sessionId)
                 throws IllegalArgumentException {
-            if ((sessionId != AudioManager.AUDIO_SESSION_ID_GENERATE) && (sessionId < 1)) {
+            if ((sessionId != AUDIO_SESSION_ID_GENERATE) && (sessionId < 1)) {
                 throw new IllegalArgumentException("Invalid audio session ID " + sessionId);
             }
             mSessionId = sessionId;
@@ -1371,8 +1402,8 @@ public class AudioTrack extends PlayerBase
 
             try {
                 final AudioTrack track = new AudioTrack(
-                        mAttributes, mFormat, mBufferSizeInBytes, mMode, mSessionId, mOffload,
-                        mEncapsulationMode, mTunerConfiguration);
+                        mContext, mAttributes, mFormat, mBufferSizeInBytes, mMode, mSessionId,
+                        mOffload, mEncapsulationMode, mTunerConfiguration);
                 if (track.getState() == STATE_UNINITIALIZED) {
                     // release is not necessary
                     throw new UnsupportedOperationException("Cannot create AudioTrack");
@@ -1384,6 +1415,32 @@ public class AudioTrack extends PlayerBase
         }
     }
 
+    /**
+     * Helper method to extract device specific audio session id from Context.
+     *
+     * @param context {@link Context} to use for extraction of device specific session id.
+     * @return device specific session id. If context is null or doesn't have specific audio
+     *   session id associated, this method returns {@link AUDIO_SESSION_ID_GENERATE}.
+     */
+    private static int getSessionIdForContext(@Nullable Context context) {
+        if (context == null) {
+            return AUDIO_SESSION_ID_GENERATE;
+        }
+
+        int deviceId = context.getDeviceId();
+        if (deviceId == DEVICE_ID_DEFAULT) {
+            return AUDIO_SESSION_ID_GENERATE;
+        }
+
+        VirtualDeviceManager vdm = context.getSystemService(VirtualDeviceManager.class);
+        if (vdm == null || vdm.getDevicePolicy(deviceId, POLICY_TYPE_AUDIO)
+                == DEVICE_POLICY_DEFAULT) {
+            return AUDIO_SESSION_ID_GENERATE;
+        }
+
+        return vdm.getAudioPlaybackSessionId(deviceId);
+    }
+
     /**
      * Sets an {@link AudioPolicy} to automatically unregister when the track is released.
      *
diff --git a/media/tests/MediaFrameworkTest/Android.bp b/media/tests/MediaFrameworkTest/Android.bp
index 48d56d898138f..06ec949e67365 100644
--- a/media/tests/MediaFrameworkTest/Android.bp
+++ b/media/tests/MediaFrameworkTest/Android.bp
@@ -13,13 +13,18 @@ android_test {
     libs: [
         "android.test.runner",
         "android.test.base",
+        "android.test.mock",
     ],
     static_libs: [
-        "mockito-target-minus-junit4",
+        "mockito-target-inline-minus-junit4",
         "androidx.test.ext.junit",
         "androidx.test.rules",
         "android-ex-camera2",
-        "testng"
+        "testng",
+    ],
+    jni_libs: [
+        "libdexmakerjvmtiagent",
+        "libstaticjvmtiagent",
     ],
     platform_apis: true,
 }
diff --git a/media/tests/MediaFrameworkTest/AndroidManifest.xml b/media/tests/MediaFrameworkTest/AndroidManifest.xml
index 9211a7e341cd6..41ac285e49f51 100644
--- a/media/tests/MediaFrameworkTest/AndroidManifest.xml
+++ b/media/tests/MediaFrameworkTest/AndroidManifest.xml
@@ -24,7 +24,8 @@
     <uses-permission android:name="android.permission.MODIFY_AUDIO_SETTINGS"/>
     <uses-permission android:name="android.permission.WAKE_LOCK"/>
 
-    <application android:networkSecurityConfig="@xml/network_security_config">
+    <application android:networkSecurityConfig="@xml/network_security_config"
+        android:debuggable="true">
         <uses-library android:name="android.test.runner"/>
         <activity android:label="@string/app_name"
              android:name="MediaFrameworkTest"
diff --git a/media/tests/MediaFrameworkTest/src/com/android/mediaframeworktest/unit/AudioTrackUnitTest.java b/media/tests/MediaFrameworkTest/src/com/android/mediaframeworktest/unit/AudioTrackUnitTest.java
new file mode 100644
index 0000000000000..ffed39ac82bf5
--- /dev/null
+++ b/media/tests/MediaFrameworkTest/src/com/android/mediaframeworktest/unit/AudioTrackUnitTest.java
@@ -0,0 +1,145 @@
+/*
+ * 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.mediaframeworktest.unit;
+
+import static android.companion.virtual.VirtualDeviceManager.DEVICE_ID_DEFAULT;
+import static android.companion.virtual.VirtualDeviceParams.DEVICE_POLICY_CUSTOM;
+import static android.companion.virtual.VirtualDeviceParams.DEVICE_POLICY_DEFAULT;
+import static android.companion.virtual.VirtualDeviceParams.POLICY_TYPE_AUDIO;
+import static android.media.AudioManager.AUDIO_SESSION_ID_GENERATE;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotEquals;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.companion.virtual.VirtualDeviceManager;
+import android.content.Context;
+import android.media.AudioAttributes;
+import android.media.AudioFormat;
+import android.media.AudioManager;
+import android.media.AudioTrack;
+import android.test.mock.MockContext;
+
+import androidx.test.ext.junit.runners.AndroidJUnit4;
+import androidx.test.platform.app.InstrumentationRegistry;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(AndroidJUnit4.class)
+public class AudioTrackUnitTest {
+    private static final int TEST_SAMPLE_RATE = 44000;
+    private static final int TEST_VIRTUAL_DEVICE_ID = 42;
+    private static final AudioAttributes TEST_AUDIO_ATTRIBUTES =
+            new AudioAttributes.Builder().setUsage(AudioAttributes.USAGE_MEDIA).build();
+    private static final AudioFormat TEST_AUDIO_FORMAT = new AudioFormat.Builder().setSampleRate(
+            TEST_SAMPLE_RATE).setEncoding(AudioFormat.ENCODING_PCM_16BIT).setChannelMask(
+            AudioFormat.CHANNEL_OUT_STEREO).build();
+    private static final int TEST_BUFFER_SIZE = AudioTrack.getMinBufferSize(
+            TEST_AUDIO_FORMAT.getSampleRate(),
+            TEST_AUDIO_FORMAT.getChannelMask(), TEST_AUDIO_FORMAT.getEncoding());
+
+    @Test
+    public void testBuilderConstructionWithContext_defaultDeviceExplicitSessionId() {
+        Context mockDefaultDeviceContext = getVirtualDeviceMockContext(
+                DEVICE_ID_DEFAULT, /*vdm=*/null);
+        int sessionId = getContext().getSystemService(AudioManager.class).generateAudioSessionId();
+
+        AudioTrack audioTrack = new AudioTrack.Builder().setContext(
+                mockDefaultDeviceContext).setAudioAttributes(TEST_AUDIO_ATTRIBUTES).setAudioFormat(
+                TEST_AUDIO_FORMAT).setBufferSizeInBytes(TEST_BUFFER_SIZE).setSessionId(
+                sessionId).build();
+
+        assertEquals(AudioTrack.STATE_INITIALIZED, audioTrack.getState());
+        assertEquals(sessionId, audioTrack.getAudioSessionId());
+    }
+
+    @Test
+    public void testBuilderConstructionWithContext_virtualDeviceDefaultAudioPolicy() {
+        int vdmPlaybackSessionId = getContext().getSystemService(
+                AudioManager.class).generateAudioSessionId();
+        VirtualDeviceManager mockVdm = getMockVirtualDeviceManager(TEST_VIRTUAL_DEVICE_ID,
+                vdmPlaybackSessionId, DEVICE_POLICY_DEFAULT);
+        Context virtualDeviceContext = getVirtualDeviceMockContext(TEST_VIRTUAL_DEVICE_ID, mockVdm);
+
+        AudioTrack audioTrack = new AudioTrack.Builder().setContext(
+                virtualDeviceContext).setAudioAttributes(TEST_AUDIO_ATTRIBUTES).setAudioFormat(
+                TEST_AUDIO_FORMAT).setBufferSizeInBytes(TEST_BUFFER_SIZE).build();
+
+        assertEquals(AudioTrack.STATE_INITIALIZED, audioTrack.getState());
+        assertNotEquals(vdmPlaybackSessionId, audioTrack.getAudioSessionId());
+    }
+
+    @Test
+    public void testBuilderConstructionWithContext_virtualDeviceCustomAudioPolicy() {
+        int vdmPlaybackSessionId = getContext().getSystemService(
+                AudioManager.class).generateAudioSessionId();
+        VirtualDeviceManager mockVdm = getMockVirtualDeviceManager(TEST_VIRTUAL_DEVICE_ID,
+                vdmPlaybackSessionId, DEVICE_POLICY_CUSTOM);
+        Context virtualDeviceContext = getVirtualDeviceMockContext(TEST_VIRTUAL_DEVICE_ID, mockVdm);
+
+        AudioTrack audioTrack = new AudioTrack.Builder().setContext(
+                virtualDeviceContext).setAudioAttributes(TEST_AUDIO_ATTRIBUTES).setAudioFormat(
+                TEST_AUDIO_FORMAT).setBufferSizeInBytes(TEST_BUFFER_SIZE).build();
+
+        assertEquals(AudioTrack.STATE_INITIALIZED, audioTrack.getState());
+        assertEquals(vdmPlaybackSessionId, audioTrack.getAudioSessionId());
+    }
+
+    @Test
+    public void testBuilderConstructionWithContext_virtualDeviceSetSessionIdOverridesContext() {
+        int vdmPlaybackSessionId = getContext().getSystemService(
+                AudioManager.class).generateAudioSessionId();
+        int anotherSessionId = getContext().getSystemService(
+                AudioManager.class).generateAudioSessionId();
+        VirtualDeviceManager mockVdm = getMockVirtualDeviceManager(TEST_VIRTUAL_DEVICE_ID,
+                vdmPlaybackSessionId, DEVICE_POLICY_CUSTOM);
+        Context virtualDeviceContext = getVirtualDeviceMockContext(TEST_VIRTUAL_DEVICE_ID, mockVdm);
+
+        AudioTrack audioTrack = new AudioTrack.Builder().setContext(
+                        virtualDeviceContext).setAudioAttributes(
+                        TEST_AUDIO_ATTRIBUTES).setAudioFormat(
+                        TEST_AUDIO_FORMAT).setBufferSizeInBytes(TEST_BUFFER_SIZE).setSessionId(
+                        anotherSessionId).build();
+
+        assertEquals(AudioTrack.STATE_INITIALIZED, audioTrack.getState());
+        assertEquals(anotherSessionId, audioTrack.getAudioSessionId());
+    }
+
+    private Context getContext() {
+        return InstrumentationRegistry.getInstrumentation().getContext();
+    }
+
+    private static Context getVirtualDeviceMockContext(int deviceId, VirtualDeviceManager vdm) {
+        MockContext mockContext = mock(MockContext.class);
+        when(mockContext.getDeviceId()).thenReturn(deviceId);
+        when(mockContext.getSystemService(VirtualDeviceManager.class)).thenReturn(vdm);
+        return mockContext;
+    }
+
+    private static VirtualDeviceManager getMockVirtualDeviceManager(int deviceId,
+            int playbackSessionId, int audioDevicePolicy) {
+        VirtualDeviceManager vdmMock = mock(VirtualDeviceManager.class);
+        when(vdmMock.getAudioPlaybackSessionId(anyInt())).thenReturn(AUDIO_SESSION_ID_GENERATE);
+        when(vdmMock.getAudioPlaybackSessionId(deviceId)).thenReturn(playbackSessionId);
+        when(vdmMock.getDevicePolicy(anyInt(), anyInt())).thenReturn(DEVICE_POLICY_DEFAULT);
+        when(vdmMock.getDevicePolicy(deviceId, POLICY_TYPE_AUDIO)).thenReturn(audioDevicePolicy);
+        return vdmMock;
+    }
+}
diff --git a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceParamsTest.java b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceParamsTest.java
index aefe4b6f4c0ac..0f9523122199f 100644
--- a/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceParamsTest.java
+++ b/services/tests/servicestests/src/com/android/server/companion/virtual/VirtualDeviceParamsTest.java
@@ -17,6 +17,7 @@
 package com.android.server.companion.virtual;
 
 import static android.companion.virtual.VirtualDeviceParams.DEVICE_POLICY_CUSTOM;
+import static android.companion.virtual.VirtualDeviceParams.POLICY_TYPE_AUDIO;
 import static android.companion.virtual.VirtualDeviceParams.POLICY_TYPE_SENSORS;
 import static android.hardware.Sensor.TYPE_ACCELEROMETER;
 
@@ -47,6 +48,7 @@ public class VirtualDeviceParamsTest {
                 .setLockState(VirtualDeviceParams.LOCK_STATE_ALWAYS_UNLOCKED)
                 .setUsersWithMatchingAccounts(Set.of(UserHandle.of(123), UserHandle.of(456)))
                 .setDevicePolicy(POLICY_TYPE_SENSORS, DEVICE_POLICY_CUSTOM)
+                .setDevicePolicy(POLICY_TYPE_AUDIO, DEVICE_POLICY_CUSTOM)
                 .addVirtualSensorConfig(
                         new VirtualSensorConfig.Builder(TYPE_ACCELEROMETER, SENSOR_NAME)
                                 .setVendor(SENSOR_VENDOR)
@@ -62,6 +64,7 @@ public class VirtualDeviceParamsTest {
         assertThat(params.getUsersWithMatchingAccounts())
                 .containsExactly(UserHandle.of(123), UserHandle.of(456));
         assertThat(params.getDevicePolicy(POLICY_TYPE_SENSORS)).isEqualTo(DEVICE_POLICY_CUSTOM);
+        assertThat(params.getDevicePolicy(POLICY_TYPE_AUDIO)).isEqualTo(DEVICE_POLICY_CUSTOM);
 
         List<VirtualSensorConfig> sensorConfigs = params.getVirtualSensorConfigs();
         assertThat(sensorConfigs).hasSize(1);