Replace class level unwrap() with unwrapAs()

This allows classes to be unwrapped as more than one type, enabling
a few new usecases (such as unwrapping both a CaptureResult as well
as the underlying CaptureSession and CameraDevice from the
ResultMetadata object) for compatibility reasons.

* Replaced unwrap() method with unwrapAs()
* Updated implementations to honor the provided type.
* Updated fake implementations to return null instead of throwing.

Test: ./gradlew\
 :camera:camera-camera2-pipe:testDebugUnitTest\
 :camera:camera-camera2-pipe-testing:testDebugUnitTest\
 :camera:camera-camera2-pipe-integration:testDebugUnitTest

Change-Id: Ie297b8debe3741e7c69ee0838d860bcbbe5ea03c
diff --git a/camera/camera-camera2-pipe-integration/src/androidTest/java/androidx/camera/camera2/pipe/integration/Camera2CameraControlDeviceTest.kt b/camera/camera-camera2-pipe-integration/src/androidTest/java/androidx/camera/camera2/pipe/integration/Camera2CameraControlDeviceTest.kt
index 34afbfe..95ed3c0 100644
--- a/camera/camera-camera2-pipe-integration/src/androidTest/java/androidx/camera/camera2/pipe/integration/Camera2CameraControlDeviceTest.kt
+++ b/camera/camera-camera2-pipe-integration/src/androidTest/java/androidx/camera/camera2/pipe/integration/Camera2CameraControlDeviceTest.kt
@@ -36,6 +36,7 @@
 import android.hardware.camera2.CaptureRequest.CONTROL_CAPTURE_INTENT_MANUAL
 import android.hardware.camera2.CaptureRequest.Key
 import android.hardware.camera2.CaptureRequest.SCALER_CROP_REGION
+import android.hardware.camera2.TotalCaptureResult
 import android.hardware.camera2.params.MeteringRectangle
 import android.os.Build
 import androidx.camera.camera2.pipe.FrameInfo
@@ -413,7 +414,7 @@
         // Assert.
         registerListener().verify(
             { _, captureResult: FrameInfo ->
-                captureResult.unwrap()!!.let { totalCaptureResult ->
+                captureResult.unwrapAs(TotalCaptureResult::class)!!.let { totalCaptureResult ->
                     if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.S) {
                         totalCaptureResult.physicalCameraTotalResults.containsKey(
                             physicalCameraId
diff --git a/camera/camera-camera2-pipe-testing/src/main/java/androidx/camera/camera2/pipe/testing/FakeMetadata.kt b/camera/camera-camera2-pipe-testing/src/main/java/androidx/camera/camera2/pipe/testing/FakeMetadata.kt
index acd475c7..3793665 100644
--- a/camera/camera-camera2-pipe-testing/src/main/java/androidx/camera/camera2/pipe/testing/FakeMetadata.kt
+++ b/camera/camera-camera2-pipe-testing/src/main/java/androidx/camera/camera2/pipe/testing/FakeMetadata.kt
@@ -36,6 +36,7 @@
 import androidx.camera.camera2.pipe.RequestTemplate
 import androidx.camera.camera2.pipe.FrameMetadata
 import androidx.camera.camera2.pipe.StreamId
+import kotlin.reflect.KClass
 import kotlinx.atomicfu.atomic
 
 private val fakeCameraIds = atomic(0)
@@ -99,12 +100,7 @@
     override fun awaitPhysicalMetadata(cameraId: CameraId): CameraMetadata =
         physicalMetadata[cameraId]!!
 
-    /** @throws UnsupportedOperationException */
-    override fun unwrap(): CameraCharacteristics? {
-        throw UnsupportedOperationException(
-            "FakeCameraMetadata does not wrap CameraCharacteristics"
-        )
-    }
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = null
 }
 
 /**
@@ -123,12 +119,7 @@
     override fun <T> get(key: CaptureRequest.Key<T>): T? = requestParameters[key] as T?
     override fun <T> getOrDefault(key: CaptureRequest.Key<T>, default: T): T = get(key) ?: default
 
-    /** @throws UnsupportedOperationException */
-    override fun unwrap(): CaptureRequest? {
-        throw UnsupportedOperationException(
-            "FakeRequestMetadata does not wrap a real CaptureRequest"
-        )
-    }
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = null
 }
 
 /**
@@ -147,12 +138,7 @@
 
     override fun <T> getOrDefault(key: CaptureResult.Key<T>, default: T): T = get(key) ?: default
 
-    /** @throws UnsupportedOperationException */
-    override fun unwrap(): CaptureResult? {
-        throw UnsupportedOperationException(
-            "FakeFrameMetadata does not wrap a real CaptureResult"
-        )
-    }
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = null
 }
 
 /**
@@ -171,11 +157,5 @@
 
     override val frameNumber: FrameNumber
         get() = metadata.frameNumber
-
-    /** @throws UnsupportedOperationException */
-    override fun unwrap(): TotalCaptureResult? {
-        throw UnsupportedOperationException(
-            "FakeFrameInfo does not wrap a real TotalCaptureResult object!"
-        )
-    }
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = null
 }
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraMetadata.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraMetadata.kt
index f3eccd4..78f1d39 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraMetadata.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraMetadata.kt
@@ -33,7 +33,7 @@
  * across all OS levels and makes behavior that depends on [CameraMetadata] easier to test and
  * reason about.
  */
-public interface CameraMetadata : Metadata, UnsafeWrapper<CameraCharacteristics> {
+public interface CameraMetadata : Metadata, UnsafeWrapper {
     public operator fun <T> get(key: CameraCharacteristics.Key<T>): T?
     public fun <T> getOrDefault(key: CameraCharacteristics.Key<T>, default: T): T
 
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/Metadata.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/Metadata.kt
index 2022d8c..2c1067c 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/Metadata.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/Metadata.kt
@@ -70,7 +70,7 @@
  * [CameraGraph]. This class will report the actual keys / values that were sent to camera2 (if
  * different) from the request that was used to create the Camera2 [CaptureRequest].
  */
-public interface RequestMetadata : Metadata, UnsafeWrapper<CaptureRequest> {
+public interface RequestMetadata : Metadata, UnsafeWrapper {
     public operator fun <T> get(key: CaptureRequest.Key<T>): T?
     public fun <T> getOrDefault(key: CaptureRequest.Key<T>, default: T): T
 
@@ -97,7 +97,7 @@
 /**
  * [FrameInfo] is a wrapper around [TotalCaptureResult].
  */
-public interface FrameInfo : UnsafeWrapper<TotalCaptureResult> {
+public interface FrameInfo : UnsafeWrapper {
     public val metadata: FrameMetadata
 
     /**
@@ -114,7 +114,7 @@
 /**
  * [FrameMetadata] is a wrapper around [CaptureResult].
  */
-public interface FrameMetadata : Metadata, UnsafeWrapper<CaptureResult> {
+public interface FrameMetadata : Metadata, UnsafeWrapper {
     public operator fun <T> get(key: CaptureResult.Key<T>): T?
     public fun <T> getOrDefault(key: CaptureResult.Key<T>, default: T): T
 
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/UnsafeWrapper.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/UnsafeWrapper.kt
index 4218951..15c221d 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/UnsafeWrapper.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/UnsafeWrapper.kt
@@ -17,14 +17,27 @@
 package androidx.camera.camera2.pipe
 
 import androidx.annotation.RequiresApi
+import kotlin.reflect.KClass
 
 /**
+ * An interface for wrapper objects that should not normally be accessed directly.
+ *
  * This interface indicates that an object or interface wraps a specific Android object or type and
  * provides a way to retrieve the underlying object directly. Accessing the underlying objects can
- * be useful for compatibility and testing, but it is extremely risky if the lifetime of the object
- * is managed by Camera Pipe and the wrapped object is closed, released, or altered.
+ * be useful for compatibility and testing, but is extremely risky if the state or lifetime of the
+ * of the object is managed by CameraPipe.
  */
 @RequiresApi(21) // TODO(b/200306659): Remove and replace with annotation on package-info.java
-public interface UnsafeWrapper<T> {
-    public fun unwrap(): T?
+public interface UnsafeWrapper {
+    /**
+     * Attempt to unwrap this object into an underlying type.
+     *
+     * This operation is not safe and should be used with caution as it makes no guarantees about
+     * the state of the underlying objects. In particular, implementations should assume that fakes,
+     * test wrappers will always return null. Finally this method should return null when unwrapping
+     * into the provided type is not supported.
+     *
+     * @return unwrapped object matching T or null
+     */
+    public fun <T : Any> unwrapAs(type: KClass<T>): T?
 }
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CameraMetadata.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CameraMetadata.kt
index 95608eb..652299c 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CameraMetadata.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CameraMetadata.kt
@@ -27,6 +27,7 @@
 import androidx.camera.camera2.pipe.CameraMetadata
 import androidx.camera.camera2.pipe.Metadata
 import androidx.camera.camera2.pipe.core.Debug
+import kotlin.reflect.KClass
 
 /**
  * This implementation provides access to [CameraCharacteristics] and lazy caching of properties
@@ -87,7 +88,11 @@
     override fun <T> getOrDefault(key: CameraCharacteristics.Key<T>, default: T): T =
         get(key) ?: default
 
-    override fun unwrap(): CameraCharacteristics = characteristics
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = when (type) {
+        CameraCharacteristics::class -> characteristics as T
+        else -> null
+    }
 
     override val keys: Set<CameraCharacteristics.Key<*>> get() = _keys.value
     override val requestKeys: Set<CaptureRequest.Key<*>> get() = _requestKeys.value
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CaptureSequenceProcessor.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CaptureSequenceProcessor.kt
index b831791..267b739 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CaptureSequenceProcessor.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Camera2CaptureSequenceProcessor.kt
@@ -36,6 +36,7 @@
 import androidx.camera.camera2.pipe.core.Threads
 import androidx.camera.camera2.pipe.writeParameters
 import javax.inject.Inject
+import kotlin.reflect.KClass
 import kotlinx.atomicfu.atomic
 
 internal interface Camera2CaptureSequenceProcessorFactory {
@@ -307,5 +308,9 @@
 
     override fun <T> getOrDefault(key: Metadata.Key<T>, default: T): T = get(key) ?: default
 
-    override fun unwrap(): CaptureRequest = captureRequest
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = when (type) {
+        CaptureRequest::class -> captureRequest as T
+        else -> null
+    }
 }
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CameraDeviceWrapper.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CameraDeviceWrapper.kt
index 464e312..9115b43 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CameraDeviceWrapper.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CameraDeviceWrapper.kt
@@ -22,6 +22,7 @@
 import android.hardware.camera2.CaptureRequest
 import android.hardware.camera2.TotalCaptureResult
 import android.hardware.camera2.params.InputConfiguration
+import android.hardware.camera2.params.OutputConfiguration
 import android.os.Build
 import android.os.Handler
 import android.view.Surface
@@ -35,6 +36,7 @@
 import androidx.camera.camera2.pipe.core.Timestamps
 import androidx.camera.camera2.pipe.core.Timestamps.formatMs
 import androidx.camera.camera2.pipe.writeParameter
+import kotlin.reflect.KClass
 import kotlinx.atomicfu.atomic
 
 /** Interface around a [CameraDevice] with minor modifications.
@@ -42,7 +44,7 @@
  * This interface has been modified to correct nullness, adjust exceptions, and to return or produce
  * wrapper interfaces instead of the native Camera2 types.
  */
-internal interface CameraDeviceWrapper : UnsafeWrapper<CameraDevice> {
+internal interface CameraDeviceWrapper : UnsafeWrapper {
     /** @see [CameraDevice.getId] */
     val cameraId: CameraId
 
@@ -112,7 +114,7 @@
 
 internal fun CameraDeviceWrapper?.closeWithTrace() {
     this?.let {
-        it.unwrap().closeWithTrace()
+        it.unwrapAs(CameraDevice::class).closeWithTrace()
         it.onDeviceClosed()
     }
 }
@@ -134,7 +136,7 @@
     private val cameraMetadata: CameraMetadata,
     private val cameraDevice: CameraDevice,
     override val cameraId: CameraId
-) : CameraDeviceWrapper, UnsafeWrapper<CameraDevice> {
+) : CameraDeviceWrapper {
     private val _lastStateCallback = atomic<CameraCaptureSessionWrapper.StateCallback?>(null)
 
     override fun createCaptureSession(
@@ -212,7 +214,7 @@
         // running on older versions of the OS.
         Api24Compat.createCaptureSessionByOutputConfigurations(
             cameraDevice,
-            outputConfigurations.map { it.unwrap() },
+            outputConfigurations.map { it.unwrapAs(OutputConfiguration::class) },
             AndroidCaptureSessionStateCallback(this, stateCallback, previousStateCallback),
             handler
         )
@@ -236,7 +238,7 @@
             Api23Compat.newInputConfiguration(
                 inputConfig.width, inputConfig.height, inputConfig.format
             ),
-            outputs.map { it.unwrap() },
+            outputs.map { it.unwrapAs(OutputConfiguration::class) },
             AndroidCaptureSessionStateCallback(this, stateCallback, previousStateCallback),
             handler
         )
@@ -250,7 +252,7 @@
 
         val sessionConfig = Api28Compat.newSessionConfiguration(
             config.sessionType,
-            config.outputConfigurations.map { it.unwrap() },
+            config.outputConfigurations.map { it.unwrapAs(OutputConfiguration::class) },
             config.executor,
             AndroidCaptureSessionStateCallback(this, stateCallback, previousStateCallback)
         )
@@ -301,7 +303,10 @@
         lastStateCallback?.onSessionFinalized()
     }
 
-    override fun unwrap(): CameraDevice? {
-        return cameraDevice
-    }
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? =
+        when (type) {
+            CameraDevice::class -> cameraDevice as T
+            else -> null
+        }
 }
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CaptureSessionWrapper.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CaptureSessionWrapper.kt
index 71dfc26..1b1ad62 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CaptureSessionWrapper.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/CaptureSessionWrapper.kt
@@ -21,6 +21,7 @@
 import android.hardware.camera2.CameraCaptureSession
 import android.hardware.camera2.CameraConstrainedHighSpeedCaptureSession
 import android.hardware.camera2.CaptureRequest
+import android.hardware.camera2.params.OutputConfiguration
 import android.os.Build
 import android.os.Handler
 import android.view.Surface
@@ -28,6 +29,7 @@
 import androidx.camera.camera2.pipe.UnsafeWrapper
 import androidx.camera.camera2.pipe.core.Log
 import java.io.Closeable
+import kotlin.reflect.KClass
 import kotlinx.atomicfu.atomic
 
 /**
@@ -36,7 +38,7 @@
  * This interface has been modified to correct nullness, adjust exceptions, and to return or produce
  * wrapper interfaces instead of the native Camera2 types.
  */
-internal interface CameraCaptureSessionWrapper : UnsafeWrapper<CameraCaptureSession>, Closeable {
+internal interface CameraCaptureSessionWrapper : UnsafeWrapper, Closeable {
 
     /**
      * @see [CameraCaptureSession.getDevice]
@@ -256,6 +258,7 @@
         finalizeLastSession()
         stateCallback.onSessionFinalized()
     }
+
     private fun finalizeLastSession() {
         // Clear out the reference to the previous session, if one was set.
         val previousSession = _lastStateCallback.getAndSet(null)
@@ -355,15 +358,15 @@
         rethrowCamera2Exceptions {
             Api26Compat.finalizeOutputConfigurations(
                 cameraCaptureSession,
-                outputConfigs.map {
-                    it.unwrap()
-                }
+                outputConfigs.map { it.unwrapAs(OutputConfiguration::class) }
             )
         }
     }
 
-    override fun unwrap(): CameraCaptureSession? {
-        return cameraCaptureSession
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = when (type) {
+        CameraCaptureSession::class -> cameraCaptureSession as T?
+        else -> null
     }
 
     override fun close() {
@@ -407,4 +410,10 @@
             throw ObjectUnavailableException(e)
         }
     }
+
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = when (type) {
+        CameraConstrainedHighSpeedCaptureSession::class -> session as T?
+        else -> super.unwrapAs(type)
+    }
 }
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Configuration.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Configuration.kt
index 3079060..a4fc62e 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Configuration.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/Configuration.kt
@@ -38,6 +38,7 @@
 import androidx.camera.camera2.pipe.core.checkOOrHigher
 import androidx.camera.camera2.pipe.core.checkPOrHigher
 import java.util.concurrent.Executor
+import kotlin.reflect.KClass
 
 /**
  * A data class that mirrors the fields in [android.hardware.camera2.params.SessionConfiguration] so
@@ -81,7 +82,7 @@
  * [OutputConfiguration]'s are NOT immutable, and changing state of an [OutputConfiguration] may
  * require the CameraCaptureSession to be finalized or updated.
  */
-internal interface OutputConfigurationWrapper : UnsafeWrapper<OutputConfiguration> {
+internal interface OutputConfigurationWrapper : UnsafeWrapper {
     /**
      * This method will return null if the output configuration was created without a Surface,
      * and until addSurface is called for the first time.
@@ -306,7 +307,11 @@
     override val surfaceGroupId: Int
         get() = output.surfaceGroupId
 
-    override fun unwrap(): OutputConfiguration = output
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = when (type) {
+        OutputConfiguration::class -> output as T
+        else -> null
+    }
 
     override fun toString(): String = output.toString()
 }
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/ExternalRequestProcessor.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/ExternalRequestProcessor.kt
index 704af63..1fcd0e4 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/ExternalRequestProcessor.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/ExternalRequestProcessor.kt
@@ -36,6 +36,7 @@
 import androidx.camera.camera2.pipe.core.Log
 import androidx.camera.camera2.pipe.graph.GraphListener
 import androidx.camera.camera2.pipe.graph.GraphRequestProcessor
+import kotlin.reflect.KClass
 import kotlinx.atomicfu.atomic
 
 @RequiresApi(21)
@@ -226,9 +227,6 @@
 
         override fun <T> getOrDefault(key: Metadata.Key<T>, default: T): T = get(key) ?: default
 
-        override fun unwrap(): CaptureRequest? {
-            // CustomRequestMetadata does not extend a Camera2 CaptureRequest.
-            return null
-        }
+        override fun <T : Any> unwrapAs(type: KClass<T>): T? = null
     }
 }
\ No newline at end of file
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/FrameMetadata.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/FrameMetadata.kt
index 0a13bd1..8a648c6 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/FrameMetadata.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/FrameMetadata.kt
@@ -29,6 +29,7 @@
 import androidx.camera.camera2.pipe.FrameNumber
 import androidx.camera.camera2.pipe.Metadata
 import androidx.camera.camera2.pipe.RequestMetadata
+import kotlin.reflect.KClass
 
 /**
  * An implementation of [FrameMetadata] that retrieves values from a [CaptureResult] object
@@ -54,7 +55,11 @@
 
     override val extraMetadata: Map<*, Any?> = emptyMap<Any, Any>()
 
-    override fun unwrap(): CaptureResult? = null
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = when (type) {
+        CaptureResult::class -> captureResult as T
+        else -> null
+    }
 }
 
 /**
@@ -83,7 +88,8 @@
     override val frameNumber: FrameNumber
         get() = frameMetadata.frameNumber
 
-    override fun unwrap(): CaptureResult? = frameMetadata.unwrap()
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = frameMetadata.unwrapAs(type)
 }
 
 /**
@@ -135,5 +141,6 @@
     override val frameNumber: FrameNumber
         get() = result.frameNumber
 
-    override fun unwrap(): TotalCaptureResult? = totalCaptureResult
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = totalCaptureResult as? T?
 }
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/VirtualCamera.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/VirtualCamera.kt
index 6bbd539..231ec8b 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/VirtualCamera.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/compat/VirtualCamera.kt
@@ -238,7 +238,7 @@
         }
 
         closeWith(
-            device?.unwrap(),
+            device?.unwrapAs(CameraDevice::class),
             @Suppress("SyntheticAccessor")
             ClosingInfo(
                 ClosedReason.APP_CLOSED
diff --git a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/compat/VirtualCameraTest.kt b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/compat/VirtualCameraTest.kt
index cca0004..4d1cfb7 100644
--- a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/compat/VirtualCameraTest.kt
+++ b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/compat/VirtualCameraTest.kt
@@ -16,6 +16,7 @@
 
 package androidx.camera.camera2.pipe.compat
 
+import android.hardware.camera2.CameraDevice
 import android.os.Build
 import android.os.Looper.getMainLooper
 import androidx.camera.camera2.pipe.core.Timestamps
@@ -185,7 +186,11 @@
         listener.onOpened(testCamera.cameraDevice)
 
         assertThat(listener.state.value).isInstanceOf(CameraStateOpen::class.java)
-        assertThat((listener.state.value as CameraStateOpen).cameraDevice.unwrap())
+        assertThat(
+            (listener.state.value as CameraStateOpen)
+                .cameraDevice
+                .unwrapAs(CameraDevice::class)
+        )
             .isSameInstanceAs(testCamera.cameraDevice)
 
         mainLooper.idleFor(1000, TimeUnit.MILLISECONDS)
diff --git a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCameraDeviceWrapper.kt b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCameraDeviceWrapper.kt
index c8ed333..c854130 100644
--- a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCameraDeviceWrapper.kt
+++ b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCameraDeviceWrapper.kt
@@ -31,6 +31,7 @@
 import androidx.camera.camera2.pipe.compat.InputConfigData
 import androidx.camera.camera2.pipe.compat.OutputConfigurationWrapper
 import androidx.camera.camera2.pipe.compat.SessionConfigData
+import kotlin.reflect.KClass
 
 /**
  * Fake implementation of [CameraDeviceWrapper] for tests.
@@ -117,5 +118,9 @@
         return nextSession
     }
 
-    override fun unwrap(): CameraDevice? = fakeCamera.cameraDevice
+    @Suppress("UNCHECKED_CAST")
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = when (type) {
+        CameraDevice::class -> fakeCamera.cameraDevice as T
+        else -> null
+    }
 }
\ No newline at end of file
diff --git a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCaptureSessionWrapper.kt b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCaptureSessionWrapper.kt
index a73f197..63ba737 100644
--- a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCaptureSessionWrapper.kt
+++ b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeCaptureSessionWrapper.kt
@@ -23,6 +23,7 @@
 import androidx.camera.camera2.pipe.compat.CameraCaptureSessionWrapper
 import androidx.camera.camera2.pipe.compat.CameraDeviceWrapper
 import androidx.camera.camera2.pipe.compat.OutputConfigurationWrapper
+import kotlin.reflect.KClass
 
 internal class FakeCaptureSessionWrapper(
     override val device: CameraDeviceWrapper,
@@ -102,12 +103,7 @@
         )
     }
 
-    override fun unwrap(): CameraCaptureSession? {
-        throw UnsupportedOperationException(
-            "FakeCaptureSessionWrapper does not wrap CameraCaptureSession"
-        )
-    }
-
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = null
     override fun close() {
         closed = true
     }
diff --git a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeOutputConfigurationWrapper.kt b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeOutputConfigurationWrapper.kt
index 8e5a267..60fa9c1 100644
--- a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeOutputConfigurationWrapper.kt
+++ b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/testing/FakeOutputConfigurationWrapper.kt
@@ -16,10 +16,10 @@
 
 package androidx.camera.camera2.pipe.testing
 
-import android.hardware.camera2.params.OutputConfiguration
 import android.view.Surface
 import androidx.camera.camera2.pipe.CameraId
 import androidx.camera.camera2.pipe.compat.OutputConfigurationWrapper
+import kotlin.reflect.KClass
 
 /**
  * Fake [OutputConfigurationWrapper] for use in tests.
@@ -53,7 +53,5 @@
         _surfaces.remove(surface)
     }
 
-    override fun unwrap(): OutputConfiguration? {
-        return null
-    }
+    override fun <T : Any> unwrapAs(type: KClass<T>): T? = null
 }
\ No newline at end of file