diff --git a/core/api/current.txt b/core/api/current.txt
index 4b3e35ee1f188..64f95774d601f 100644
--- a/core/api/current.txt
+++ b/core/api/current.txt
@@ -21197,6 +21197,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 50c3c78ce9d44..567a570309a31 100644
--- a/core/api/system-current.txt
+++ b/core/api/system-current.txt
@@ -3043,6 +3043,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 cc452e27a650c..f851c319580ef 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;
@@ -254,6 +256,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);