Refactor TorchControl to simplify and optimize the logic

  The primary objective in this CL is setting AE mode at the same capture request where torch mode is set. Currently, we set the torch mode first and await its submission before updating the AE mode in a separate request. This can be error-prone in future changes as well for cases like when user sets some AE mode in the middle of these two operations. Ideally, we should be able to do a single quick synchronous update to CameraPipe when user calls enableTorch or disableTorch.

1. Allow users to override the AE mode set at Controller3A#setTorch
- This allows us to set the appropriate AE mode at the same request of disabling torch. Otherwise, two separate requests need to be used.
2. Simplify TorchControl by awaiting on Controller3A.setTorch only instead of requiring to depend on update3A as well to workaround the previous limitation.

  This can be helpful in aosp/3210473 as well, see the difference in TorchControl vs other control classes there in patch 17.

Bug: 358093572
Test: TorchControlTest, CameraControlAdapterDeviceTest
Change-Id: I54098ff833f915976d3c445c6f53b489c6176f32
diff --git a/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/State3AControl.kt b/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/State3AControl.kt
index da6fcdb..fa9fdb7 100644
--- a/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/State3AControl.kt
+++ b/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/State3AControl.kt
@@ -90,6 +90,14 @@
     public var flashMode: Int by updateOnPropertyChange(DEFAULT_FLASH_MODE)
     public var template: Int by updateOnPropertyChange(DEFAULT_REQUEST_TEMPLATE)
     public var tryExternalFlashAeMode: Boolean by updateOnPropertyChange(false)
+
+    /**
+     * The [CaptureRequest.CONTROL_AE_MODE] that is set to camera if supported.
+     *
+     * If null, a value based on other settings is calculated and available via
+     * [getFinalPreferredAeMode]. If not supported, [getSupportedAeMode] is used to find the next
+     * best option.
+     */
     public var preferredAeMode: Int? by updateOnPropertyChange(null)
     public var preferredFocusMode: Int? by updateOnPropertyChange(null)
     public var preferredAeFpsRange: Range<Int>? by
@@ -118,6 +126,19 @@
             }
         }
 
+    /**
+     * Returns the AE mode that is finally set to camera based on all other settings and camera
+     * capabilities.
+     */
+    public fun getFinalSupportedAeMode(): Int =
+        cameraProperties.metadata.getSupportedAeMode(getFinalPreferredAeMode())
+
+    /**
+     * Returns the AE mode that is finally set to camera based on all other settings.
+     *
+     * Note that this may not be supported via the camera and should be sanitized with
+     * [getSupportedAeMode].
+     */
     private fun getFinalPreferredAeMode(): Int {
         var preferAeMode =
             preferredAeMode
diff --git a/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/TorchControl.kt b/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/TorchControl.kt
index 7d2deb3..75cdcfe 100644
--- a/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/TorchControl.kt
+++ b/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/TorchControl.kt
@@ -17,6 +17,9 @@
 package androidx.camera.camera2.pipe.integration.impl
 
 import android.hardware.camera2.CaptureRequest
+import androidx.camera.camera2.pipe.AeMode
+import androidx.camera.camera2.pipe.core.Log.debug
+import androidx.camera.camera2.pipe.core.Log.warn
 import androidx.camera.camera2.pipe.integration.adapter.propagateTo
 import androidx.camera.camera2.pipe.integration.compat.workaround.isFlashAvailable
 import androidx.camera.camera2.pipe.integration.config.CameraScope
@@ -86,6 +89,8 @@
         cancelPreviousTask: Boolean = true,
         ignoreFlashUnitAvailability: Boolean = false
     ): Deferred<Unit> {
+        debug { "TorchControl#setTorchAsync: torch = $torch" }
+
         val signal = CompletableDeferred<Unit>()
 
         if (!ignoreFlashUnitAvailability && !hasFlashUnit) {
@@ -107,16 +112,31 @@
 
                 _updateSignal = signal
 
-                // TODO(b/209757083), handle the failed result of the setTorchAsync().
-                requestControl.setTorchAsync(torch).join()
-
-                // Hold the internal AE mode to ON while the torch is turned ON.
+                // Hold the internal AE mode to ON while the torch is turned ON. If torch is OFF, a
+                // value of null will make the state3AControl calculate the correct AE mode based on
+                // other settings.
                 state3AControl.preferredAeMode =
                     if (torch) CaptureRequest.CONTROL_AE_MODE_ON else null
+                val aeMode: AeMode =
+                    AeMode.fromIntOrNull(state3AControl.getFinalSupportedAeMode())
+                        ?: run {
+                            warn {
+                                "TorchControl#setTorchAsync: Failed to convert ae mode of value" +
+                                    " ${state3AControl.getFinalSupportedAeMode()} with" +
+                                    " AeMode.fromIntOrNull, fallback to AeMode.ON"
+                            }
+                            AeMode.ON
+                        }
 
-                // Always update3A again to reset the AE state in the Camera-pipe controller.
-                state3AControl.invalidate()
-                state3AControl.updateSignal?.propagateTo(signal) ?: run { signal.complete(Unit) }
+                val deferred =
+                    if (torch) requestControl.setTorchOnAsync()
+                    else requestControl.setTorchOffAsync(aeMode)
+                deferred.propagateTo(signal) {
+                    // TODO: b/209757083 - handle the failed result of the setTorchAsync().
+                    //   Since we are not handling the result here, signal is completed with Unit
+                    //   value here without exception when source deferred completes (returning Unit
+                    //   explicitly is redundant and thus this block looks empty)
+                }
             }
         }
             ?: run {
diff --git a/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/UseCaseCameraRequestControl.kt b/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/UseCaseCameraRequestControl.kt
index 0731bda..c105c916 100644
--- a/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/UseCaseCameraRequestControl.kt
+++ b/camera/camera-camera2-pipe-integration/src/main/java/androidx/camera/camera2/pipe/integration/impl/UseCaseCameraRequestControl.kt
@@ -28,7 +28,6 @@
 import androidx.camera.camera2.pipe.RequestTemplate
 import androidx.camera.camera2.pipe.Result3A
 import androidx.camera.camera2.pipe.StreamId
-import androidx.camera.camera2.pipe.TorchState
 import androidx.camera.camera2.pipe.core.Log.debug
 import androidx.camera.camera2.pipe.integration.config.UseCaseCameraScope
 import androidx.camera.camera2.pipe.integration.config.UseCaseGraphConfig
@@ -121,12 +120,20 @@
 
     // 3A
     /**
-     * Asynchronously sets the torch (flashlight) state.
+     * Asynchronously sets the torch (flashlight) to ON state.
      *
-     * @param enabled True to enable the torch, false to disable it.
      * @return A [Deferred] representing the asynchronous operation and its result ([Result3A]).
      */
-    public suspend fun setTorchAsync(enabled: Boolean): Deferred<Result3A>
+    public suspend fun setTorchOnAsync(): Deferred<Result3A>
+
+    /**
+     * Asynchronously sets the torch (flashlight) state to OFF state.
+     *
+     * @param aeMode The [AeMode] to set while setting the torch value. See
+     *   [CameraGraph.Session.setTorchOff] for details.
+     * @return A [Deferred] representing the asynchronous operation and its result ([Result3A]).
+     */
+    public suspend fun setTorchOffAsync(aeMode: AeMode): Deferred<Result3A>
 
     /**
      * Asynchronously starts a 3A (Auto Exposure, Auto Focus, Auto White Balance) operation with the
@@ -264,14 +271,14 @@
                 )
         } ?: canceledResult
 
-    override suspend fun setTorchAsync(enabled: Boolean): Deferred<Result3A> =
+    override suspend fun setTorchOnAsync(): Deferred<Result3A> =
+        runIfNotClosed { useGraphSessionOrFailed { it.setTorchOn() } } ?: submitFailedResult
+
+    override suspend fun setTorchOffAsync(aeMode: AeMode): Deferred<Result3A> =
         runIfNotClosed {
             useGraphSessionOrFailed {
-                it.setTorch(
-                    when (enabled) {
-                        true -> TorchState.ON
-                        false -> TorchState.OFF
-                    }
+                it.setTorchOff(
+                    aeMode = aeMode,
                 )
             }
         } ?: submitFailedResult
diff --git a/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/impl/CapturePipelineTest.kt b/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/impl/CapturePipelineTest.kt
index f968130..5dffedb 100644
--- a/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/impl/CapturePipelineTest.kt
+++ b/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/impl/CapturePipelineTest.kt
@@ -136,8 +136,14 @@
             val torchUpdateEventList = mutableListOf<Boolean>()
             val setTorchSemaphore = Semaphore(0)
 
-            override suspend fun setTorchAsync(enabled: Boolean): Deferred<Result3A> {
-                torchUpdateEventList.add(enabled)
+            override suspend fun setTorchOnAsync(): Deferred<Result3A> {
+                torchUpdateEventList.add(true)
+                setTorchSemaphore.release()
+                return CompletableDeferred(Result3A(Result3A.Status.OK))
+            }
+
+            override suspend fun setTorchOffAsync(aeMode: AeMode): Deferred<Result3A> {
+                torchUpdateEventList.add(false)
                 setTorchSemaphore.release()
                 return CompletableDeferred(Result3A(Result3A.Status.OK))
             }
diff --git a/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeCameraGraphSession.kt b/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeCameraGraphSession.kt
index f3fe9a8..e7a993a 100644
--- a/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeCameraGraphSession.kt
+++ b/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeCameraGraphSession.kt
@@ -29,7 +29,6 @@
 import androidx.camera.camera2.pipe.OutputStatus
 import androidx.camera.camera2.pipe.Request
 import androidx.camera.camera2.pipe.Result3A
-import androidx.camera.camera2.pipe.TorchState
 import androidx.camera.camera2.pipe.integration.testing.FakeCameraGraphSession.RequestStatus.ABORTED
 import androidx.camera.camera2.pipe.integration.testing.FakeCameraGraphSession.RequestStatus.FAILED
 import androidx.camera.camera2.pipe.integration.testing.FakeCameraGraphSession.RequestStatus.TOTAL_CAPTURE_DONE
@@ -104,7 +103,11 @@
         throw NotImplementedError("Not used in testing")
     }
 
-    override fun setTorch(torchState: TorchState): Deferred<Result3A> {
+    override fun setTorchOn(): Deferred<Result3A> {
+        throw NotImplementedError("Not used in testing")
+    }
+
+    override fun setTorchOff(aeMode: AeMode?): Deferred<Result3A> {
         throw NotImplementedError("Not used in testing")
     }
 
diff --git a/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeUseCaseCamera.kt b/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeUseCaseCamera.kt
index 9da279b..0c3df52 100644
--- a/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeUseCaseCamera.kt
+++ b/camera/camera-camera2-pipe-integration/src/test/java/androidx/camera/camera2/pipe/integration/testing/FakeUseCaseCamera.kt
@@ -123,7 +123,11 @@
         return CompletableDeferred(Unit)
     }
 
-    override suspend fun setTorchAsync(enabled: Boolean): Deferred<Result3A> {
+    override suspend fun setTorchOnAsync(): Deferred<Result3A> {
+        return setTorchResult
+    }
+
+    override suspend fun setTorchOffAsync(aeMode: AeMode): Deferred<Result3A> {
         return setTorchResult
     }
 
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraControls.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraControls.kt
index 2d032b7..869f124 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraControls.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraControls.kt
@@ -77,6 +77,18 @@
 
         @JvmStatic
         public fun fromIntOrNull(value: Int): AeMode? = values.firstOrNull { it.value == value }
+
+        @JvmStatic
+        public fun fromInt(value: Int): AeMode =
+            when (value) {
+                OFF.value -> OFF
+                ON.value -> ON
+                ON_AUTO_FLASH.value -> ON_AUTO_FLASH
+                ON_ALWAYS_FLASH.value -> ON_ALWAYS_FLASH
+                ON_AUTO_FLASH_REDEYE.value -> ON_AUTO_FLASH_REDEYE
+                ON_EXTERNAL_FLASH.value -> ON_EXTERNAL_FLASH
+                else -> ON
+            }
     }
 }
 
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraGraph.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraGraph.kt
index e27775d..4e8a4ee 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraGraph.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/CameraGraph.kt
@@ -503,20 +503,30 @@
         ): Deferred<Result3A>
 
         /**
-         * Turns the torch to ON or OFF.
+         * Turns the torch to ON.
          *
          * This method has a side effect on the currently set AE mode. Ref:
          * https://developer.android.com/reference/android/hardware/camera2/CaptureRequest#FLASH_MODE
          * To use the flash control, AE mode must be set to ON or OFF. So if the AE mode is already
          * not either ON or OFF, we will need to update the AE mode to one of those states, here we
          * will choose ON. It is the responsibility of the application layer above CameraPipe to
-         * restore the AE mode after the torch control has been used. The [update3A] method can be
-         * used to restore the AE state to a previous value.
+         * restore the AE mode after the torch control has been used. The [setTorchOff] or
+         * [update3A] method can be used to restore the AE state to a previous value.
          *
          * @return the FrameNumber at which the turn was fully turned on if switch was ON, or the
          *   FrameNumber at which it was completely turned off when the switch was OFF.
          */
-        public fun setTorch(torchState: TorchState): Deferred<Result3A>
+        public fun setTorchOn(): Deferred<Result3A>
+
+        /**
+         * Turns the torch to OFF.
+         *
+         * @param aeMode The [AeMode] to set while disabling the torch value. If null which is the
+         *   default value, the current AE mode is used.
+         * @return the FrameNumber at which the turn was fully turned on if switch was ON, or the
+         *   FrameNumber at which it was completely turned off when the switch was OFF.
+         */
+        public fun setTorchOff(aeMode: AeMode? = null): Deferred<Result3A>
 
         /**
          * Locks the auto-exposure, auto-focus and auto-whitebalance as per the given desired
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/CameraGraphSessionImpl.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/CameraGraphSessionImpl.kt
index c552afc..e89dcf6 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/CameraGraphSessionImpl.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/CameraGraphSessionImpl.kt
@@ -26,7 +26,6 @@
 import androidx.camera.camera2.pipe.Lock3ABehavior
 import androidx.camera.camera2.pipe.Request
 import androidx.camera.camera2.pipe.Result3A
-import androidx.camera.camera2.pipe.TorchState
 import androidx.camera.camera2.pipe.core.Token
 import androidx.camera.camera2.pipe.internal.FrameCaptureQueue
 import kotlinx.atomicfu.atomic
@@ -115,11 +114,16 @@
         return controller3A.submit3A(aeMode, afMode, awbMode, aeRegions, afRegions, awbRegions)
     }
 
-    override fun setTorch(torchState: TorchState): Deferred<Result3A> {
-        check(!token.released) { "Cannot call setTorch on $this after close." }
+    override fun setTorchOn(): Deferred<Result3A> {
+        check(!token.released) { "Cannot call setTorchOn on $this after close." }
         // TODO(sushilnath): First check whether the camera device has a flash unit. Ref:
         // https://developer.android.com/reference/android/hardware/camera2/CameraCharacteristics#FLASH_INFO_AVAILABLE
-        return controller3A.setTorch(torchState)
+        return controller3A.setTorchOn()
+    }
+
+    override fun setTorchOff(aeMode: AeMode?): Deferred<Result3A> {
+        check(!token.released) { "Cannot call setTorchOff on $this after close." }
+        return controller3A.setTorchOff(aeMode)
     }
 
     override suspend fun lock3A(
diff --git a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/Controller3A.kt b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/Controller3A.kt
index 177ce94..8c3c597 100644
--- a/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/Controller3A.kt
+++ b/camera/camera-camera2-pipe/src/main/java/androidx/camera/camera2/pipe/graph/Controller3A.kt
@@ -43,7 +43,6 @@
 import androidx.camera.camera2.pipe.Lock3ABehavior
 import androidx.camera.camera2.pipe.Result3A
 import androidx.camera.camera2.pipe.Result3A.Status
-import androidx.camera.camera2.pipe.TorchState
 import androidx.camera.camera2.pipe.core.Log.debug
 import kotlinx.coroutines.CompletableDeferred
 import kotlinx.coroutines.Deferred
@@ -646,14 +645,23 @@
         return listener.result
     }
 
-    fun setTorch(torchState: TorchState): Deferred<Result3A> {
-        // Determine the flash mode based on the torch state.
-        val flashMode = if (torchState == TorchState.ON) FlashMode.TORCH else FlashMode.OFF
-        // To use the flash control, AE mode must be set to ON or OFF.
+    /**
+     * Enables the torch which may require changing the [AeMode].
+     *
+     * To use [FlashMode.TORCH], either [AeMode.ON] or [AeMode.OFF] needs to be used, otherwise, the
+     * flash mode is a no-op. If the current AE mode is neither of them, this function changes the
+     * AE mode to [AeMode.ON] in order to enable the torch.
+     */
+    fun setTorchOn(): Deferred<Result3A> {
         val currAeMode = graphState3A.aeMode
         val desiredAeMode =
             if (currAeMode == AeMode.ON || currAeMode == AeMode.OFF) null else AeMode.ON
-        return update3A(aeMode = desiredAeMode, flashMode = flashMode)
+        return update3A(aeMode = desiredAeMode, flashMode = FlashMode.TORCH)
+    }
+
+    /** Disables the torch and sets a new AE mode if provided. */
+    fun setTorchOff(aeMode: AeMode? = null): Deferred<Result3A> {
+        return update3A(aeMode = aeMode, flashMode = FlashMode.OFF)
     }
 
     private fun lock3ANow(
diff --git a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/graph/Controller3ASetTorchTest.kt b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/graph/Controller3ASetTorchTest.kt
index 7e4ff9a..5490270 100644
--- a/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/graph/Controller3ASetTorchTest.kt
+++ b/camera/camera-camera2-pipe/src/test/java/androidx/camera/camera2/pipe/graph/Controller3ASetTorchTest.kt
@@ -17,6 +17,7 @@
 package androidx.camera.camera2.pipe.graph
 
 import android.hardware.camera2.CaptureRequest
+import android.hardware.camera2.CaptureRequest.CONTROL_AE_MODE_ON_AUTO_FLASH
 import android.hardware.camera2.CaptureResult
 import android.os.Build
 import androidx.camera.camera2.pipe.AeMode
@@ -24,14 +25,12 @@
 import androidx.camera.camera2.pipe.FrameNumber
 import androidx.camera.camera2.pipe.RequestNumber
 import androidx.camera.camera2.pipe.Result3A
-import androidx.camera.camera2.pipe.TorchState
 import androidx.camera.camera2.pipe.testing.FakeCameraMetadata
 import androidx.camera.camera2.pipe.testing.FakeFrameMetadata
 import androidx.camera.camera2.pipe.testing.FakeGraphProcessor
 import androidx.camera.camera2.pipe.testing.FakeRequestMetadata
 import androidx.camera.camera2.pipe.testing.RobolectricCameraPipeTestRunner
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
 import org.junit.After
@@ -39,7 +38,6 @@
 import org.junit.runner.RunWith
 import org.robolectric.annotation.Config
 
-@OptIn(ExperimentalCoroutinesApi::class)
 @RunWith(RobolectricCameraPipeTestRunner::class)
 @Config(minSdk = Build.VERSION_CODES.LOLLIPOP)
 internal class Controller3ASetTorchTest {
@@ -56,7 +54,7 @@
     }
 
     @Test
-    fun testSetTorchFailsImmediatelyWithoutRepeatingRequest() = runTest {
+    fun setTorchOn_withoutRepeatingRequest_failsImmediatelyWithNoGraphStateChange() = runTest {
         val graphProcessor2 = FakeGraphProcessor()
         val controller3A =
             Controller3A(
@@ -65,17 +63,42 @@
                 graphProcessor2.graphState3A,
                 listener3A
             )
-        val result = controller3A.setTorch(TorchState.ON)
+        val result = controller3A.setTorchOn()
         assertThat(result.await().status).isEqualTo(Result3A.Status.SUBMIT_FAILED)
         assertThat(graphProcessor2.graphState3A.flashMode).isEqualTo(FlashMode.TORCH)
     }
 
     @Test
-    fun testSetTorchOn() = runTest {
-        val result = controller3A.setTorch(TorchState.ON)
+    fun setTorchOff_withoutRepeatingRequest_failsImmediatelyWithNoGraphStateChange() = runTest {
+        val graphProcessor2 = FakeGraphProcessor()
+        val controller3A =
+            Controller3A(
+                graphProcessor2,
+                FakeCameraMetadata(),
+                graphProcessor2.graphState3A,
+                listener3A
+            )
+        val result = controller3A.setTorchOff()
+        assertThat(result.await().status).isEqualTo(Result3A.Status.SUBMIT_FAILED)
+        assertThat(graphProcessor2.graphState3A.flashMode).isEqualTo(FlashMode.OFF)
+    }
+
+    @Test
+    fun setTorchOn_updatesGraphStateWithAeModeOnAndFlashModeTorch() = runTest {
+        controller3A.setTorchOn()
         assertThat(graphState3A.aeMode!!.value).isEqualTo(CaptureRequest.CONTROL_AE_MODE_ON)
         assertThat(graphState3A.flashMode!!.value).isEqualTo(CaptureRequest.FLASH_MODE_TORCH)
+    }
+
+    @Test
+    fun setTorchOn_noCaptureResultProvided_resultIncomplete() = runTest {
+        val result = controller3A.setTorchOn()
         assertThat(result.isCompleted).isFalse()
+    }
+
+    @Test
+    fun setTorchOn_captureResultProvidedWithAeOnAndFlashTorch_returnsOkResult() = runTest {
+        val result = controller3A.setTorchOn()
 
         launch {
             listener3A.onRequestSequenceCreated(
@@ -94,17 +117,62 @@
                 )
             )
         }
+
         val result3A = result.await()
         assertThat(result3A.frameMetadata!!.frameNumber.value).isEqualTo(101L)
         assertThat(result3A.status).isEqualTo(Result3A.Status.OK)
     }
 
     @Test
-    fun testSetTorchOff() = runTest {
-        val result = controller3A.setTorch(TorchState.OFF)
-        assertThat(graphState3A.aeMode!!.value).isEqualTo(CaptureRequest.CONTROL_AE_MODE_ON)
+    fun setTorchOn_captureResultProvidedWithOnlyFlashTorch_resultIncomplete() = runTest {
+        val result = controller3A.setTorchOn()
+
+        launch {
+                listener3A.onRequestSequenceCreated(
+                    FakeRequestMetadata(requestNumber = RequestNumber(1))
+                )
+                listener3A.onPartialCaptureResult(
+                    FakeRequestMetadata(requestNumber = RequestNumber(1)),
+                    FrameNumber(101L),
+                    FakeFrameMetadata(
+                        frameNumber = FrameNumber(101L),
+                        resultMetadata =
+                            mapOf(CaptureResult.FLASH_MODE to CaptureResult.FLASH_MODE_TORCH)
+                    )
+                )
+            }
+            .join()
+
+        assertThat(result.isCompleted).isFalse()
+    }
+
+    @Test
+    fun setTorchOff_updatesGraphStateWithFlashModeOff() = runTest {
+        controller3A.setTorchOff()
         assertThat(graphState3A.flashMode!!.value).isEqualTo(CaptureRequest.FLASH_MODE_OFF)
+    }
+
+    @Test
+    fun setTorchOffWithoutAeMode_graphStateAeModeStaysNull() = runTest {
+        controller3A.setTorchOff()
+        assertThat(graphState3A.aeMode?.value).isNull() // null is default value here
+    }
+
+    @Test
+    fun setTorchOffWithAutoFlashAeMode_graphStateAeModeUpdatedToAutoFlash() = runTest {
+        controller3A.setTorchOff(aeMode = AeMode.ON_AUTO_FLASH)
+        assertThat(graphState3A.aeMode?.value).isEqualTo(CONTROL_AE_MODE_ON_AUTO_FLASH)
+    }
+
+    @Test
+    fun setTorchOff_noCaptureResultWithUpdatedStates_resultIncomplete() = runTest {
+        val result = controller3A.setTorchOff()
         assertThat(result.isCompleted).isFalse()
+    }
+
+    @Test
+    fun setTorchOffWithoutAeMode_captureResultProvidedWithFlashOff_returnsOkResult() = runTest {
+        val result = controller3A.setTorchOff()
 
         launch {
             listener3A.onRequestSequenceCreated(
@@ -115,11 +183,7 @@
                 FrameNumber(101L),
                 FakeFrameMetadata(
                     frameNumber = FrameNumber(101L),
-                    resultMetadata =
-                        mapOf(
-                            CaptureResult.CONTROL_AE_MODE to CaptureResult.CONTROL_AE_MODE_ON,
-                            CaptureResult.FLASH_MODE to CaptureResult.FLASH_MODE_OFF
-                        )
+                    resultMetadata = mapOf(CaptureResult.FLASH_MODE to CaptureResult.FLASH_MODE_OFF)
                 )
             )
         }
@@ -129,33 +193,90 @@
     }
 
     @Test
-    fun testSetTorchDoesNotChangeAeModeIfNotNeeded() = runTest {
-        graphState3A.update(aeMode = AeMode.OFF)
+    fun setTorchOffWithAutoFlashAe_captureResultProvidedWithOnlyFlashOff_resultIncomplete() =
+        runTest {
+            val result = controller3A.setTorchOff(aeMode = AeMode.ON_AUTO_FLASH)
 
-        val result = controller3A.setTorch(TorchState.ON)
-        assertThat(graphState3A.aeMode!!.value).isEqualTo(CaptureRequest.CONTROL_AE_MODE_OFF)
-        assertThat(graphState3A.flashMode!!.value).isEqualTo(CaptureRequest.FLASH_MODE_TORCH)
-        assertThat(result.isCompleted).isFalse()
-
-        launch {
-            listener3A.onRequestSequenceCreated(
-                FakeRequestMetadata(requestNumber = RequestNumber(1))
-            )
-            listener3A.onPartialCaptureResult(
-                FakeRequestMetadata(requestNumber = RequestNumber(1)),
-                FrameNumber(101L),
-                FakeFrameMetadata(
-                    frameNumber = FrameNumber(101L),
-                    resultMetadata =
-                        mapOf(
-                            CaptureResult.CONTROL_AE_MODE to CaptureResult.CONTROL_AE_MODE_OFF,
-                            CaptureResult.FLASH_MODE to CaptureResult.FLASH_MODE_TORCH
+            launch {
+                    listener3A.onRequestSequenceCreated(
+                        FakeRequestMetadata(requestNumber = RequestNumber(1))
+                    )
+                    listener3A.onPartialCaptureResult(
+                        FakeRequestMetadata(requestNumber = RequestNumber(1)),
+                        FrameNumber(101L),
+                        FakeFrameMetadata(
+                            frameNumber = FrameNumber(101L),
+                            resultMetadata =
+                                mapOf(CaptureResult.FLASH_MODE to CaptureResult.FLASH_MODE_OFF)
                         )
-                )
-            )
+                    )
+                }
+                .join()
+
+            assertThat(result.isCompleted).isFalse()
         }
-        val result3A = result.await()
-        assertThat(result3A.frameMetadata!!.frameNumber.value).isEqualTo(101L)
-        assertThat(result3A.status).isEqualTo(Result3A.Status.OK)
-    }
+
+    @Test
+    fun setTorchOffWithAutoFlashAe_captureResultProvidedWithAutoAeAndFlashOff_returnsOkResult() =
+        runTest {
+            val result = controller3A.setTorchOff(aeMode = AeMode.ON_AUTO_FLASH)
+
+            launch {
+                listener3A.onRequestSequenceCreated(
+                    FakeRequestMetadata(requestNumber = RequestNumber(1))
+                )
+                listener3A.onPartialCaptureResult(
+                    FakeRequestMetadata(requestNumber = RequestNumber(1)),
+                    FrameNumber(101L),
+                    FakeFrameMetadata(
+                        frameNumber = FrameNumber(101L),
+                        resultMetadata =
+                            mapOf(
+                                CaptureResult.CONTROL_AE_MODE to
+                                    CaptureResult.CONTROL_AE_MODE_ON_AUTO_FLASH,
+                                CaptureResult.FLASH_MODE to CaptureResult.FLASH_MODE_OFF
+                            )
+                    )
+                )
+            }
+            val result3A = result.await()
+            assertThat(result3A.frameMetadata!!.frameNumber.value).isEqualTo(101L)
+            assertThat(result3A.status).isEqualTo(Result3A.Status.OK)
+        }
+
+    @Test
+    fun setTorchOn_graphStateAlreadyAeOffSoNoChangeNeeded_aeModeUnchangedButFlashChangedToTorch() =
+        runTest {
+            graphState3A.update(aeMode = AeMode.OFF)
+
+            controller3A.setTorchOn()
+            assertThat(graphState3A.aeMode!!.value).isEqualTo(CaptureRequest.CONTROL_AE_MODE_OFF)
+            assertThat(graphState3A.flashMode!!.value).isEqualTo(CaptureRequest.FLASH_MODE_TORCH)
+        }
+
+    @Test
+    fun setTorchOnWithGraphStateAlreadyAeOff_captureResultProvidedWithFlashTorch_returnsOkResult() =
+        runTest {
+            graphState3A.update(aeMode = AeMode.OFF)
+
+            val result = controller3A.setTorchOn()
+
+            launch {
+                listener3A.onRequestSequenceCreated(
+                    FakeRequestMetadata(requestNumber = RequestNumber(1))
+                )
+                listener3A.onPartialCaptureResult(
+                    FakeRequestMetadata(requestNumber = RequestNumber(1)),
+                    FrameNumber(101L),
+                    FakeFrameMetadata(
+                        frameNumber = FrameNumber(101L),
+                        resultMetadata =
+                            mapOf(CaptureResult.FLASH_MODE to CaptureResult.FLASH_MODE_TORCH)
+                    )
+                )
+            }
+            val result3A = result.await()
+            assertThat(result3A.frameMetadata!!.frameNumber.value).isEqualTo(101L)
+            assertThat(result3A.status).isEqualTo(Result3A.Status.OK)
+        }
 }