Ask user of quota usage before making a reservation

Show a dialog that explains quota usage and billing charges when user
tries to reserve a device.

Fixes: 324076707
Test: Added
Change-Id: I1c3ca6a980a92b8cff066e017149a9a01c34ffca

6 files changed, 242 insertions, 1 deletions

diff --git a/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceTemplate.kt b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceTemplate.kt
index 7aeb00d..12479cb 100644
--- a/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceTemplate.kt
+++ b/directaccess/src/com/google/gct/directaccess/provisioner/DirectAccessDeviceTemplate.kt
@@ -31,13 +31,18 @@ import com.android.tools.adbbridge.Reservation
 import com.android.tools.idea.concurrency.AndroidDispatchers
 import com.android.tools.idea.concurrency.createChildScope
 import com.android.tools.idea.devicemanager.DeviceType
+import com.android.tools.idea.flags.StudioFlags
 import com.google.gct.directaccess.analytics.DirectAccessUsageTracker
 import com.google.gct.directaccess.directAccessCloudProjectManager
 import com.google.services.firebase.directaccess.client.findOrCreateReservation
 import com.google.services.firebase.directaccess.client.waitUntilActive
 import com.google.wireless.android.sdk.stats.DirectAccessUsageEvent.FailureReason
+import com.intellij.ide.util.PropertiesComponent
 import com.intellij.openapi.project.Project
+import com.intellij.openapi.ui.DoNotAskOption
+import com.intellij.openapi.ui.MessageDialogBuilder
 import com.intellij.openapi.ui.Messages
+import com.intellij.openapi.ui.messages.MessageDialog
 import icons.StudioIcons
 import java.time.Duration
 import javax.swing.Icon
@@ -58,6 +63,14 @@ import kotlinx.coroutines.withContext
 val SHORT_AWAITING_RESERVATION_READY_TIME_LIMIT: Duration = Duration.ofMinutes(1)
 val LONG_AWAITING_RESERVATION_READY_TIME_LIMIT: Duration = Duration.ofMinutes(15)
 
+internal const val UNKNOWN_DEVICE_DO_NOT_ASK = "device.streaming.unknown.do.not.ask"
+internal const val SPARK_SINGLE_DEVICE_DO_NOT_ASK =
+  "device.streaming.spark.single.device.do.not.ask"
+internal const val SPARK_MULTI_DEVICE_DO_NOT_ASK = "device.streaming.spark.multi.device.do.not.ask"
+internal const val BLAZE_SINGLE_DEVICE_DO_NOT_ASK =
+  "device.streaming.blaze.single.device.do.not.ask"
+internal const val BLAZE_MULTI_DEVICE_DO_NOT_ASK = "device.streaming.blaze.multi.device.do.not.ask"
+
 class DirectAccessDeviceTemplate(
   private val project: Project,
   private val deviceInfoFlow: StateFlow<DeviceInfo>,
@@ -140,6 +153,10 @@ class DirectAccessDeviceTemplate(
           throw DeviceActionDisabledException(this)
         }
         confirmWaitingTime()
+        if (!confirmUsageMinutes()) {
+          isActivationStarted.value = false
+          throw CancellationException("User declined quota usage")
+        }
 
         val reservationName =
           try {
@@ -173,6 +190,83 @@ class DirectAccessDeviceTemplate(
         }
       }
 
+      private suspend fun confirmUsageMinutes(): Boolean {
+        // Don't prompt if monthly billing is not enabled
+        if (!StudioFlags.DIRECT_ACCESS_MONTHLY_QUOTA.get()) {
+          return true
+        }
+        val billingStatus = project.directAccessCloudProjectManager?.isBillingEnabledFlow?.value
+        val isMultiDevice = devices.value.filterIsInstance<DirectAccessDeviceHandle>().isNotEmpty()
+        val persistenceKey = getPersistenceKeyForDoNoAsk(billingStatus, isMultiDevice)
+        val (title, message) = getDialogTitleAndMessage(billingStatus)
+
+        return PropertiesComponent.getInstance(project).getBoolean(persistenceKey, false) ||
+          withContext(AndroidDispatchers.uiThread) {
+            MessageDialogBuilder.yesNo(title, message)
+              .doNotAsk(
+                object : DoNotAskOption.Adapter() {
+                  override fun rememberChoice(isSelected: Boolean, exitCode: Int) {
+                    if (exitCode == MessageDialog.OK_EXIT_CODE) {
+                      PropertiesComponent.getInstance(project).setValue(persistenceKey, isSelected)
+                    }
+                  }
+                }
+              )
+              .ask(project)
+          }
+      }
+
+      private fun getPersistenceKeyForDoNoAsk(billingStatus: Boolean?, isMultiDevice: Boolean) =
+        when (billingStatus) {
+          null -> UNKNOWN_DEVICE_DO_NOT_ASK
+          false ->
+            if (isMultiDevice) SPARK_MULTI_DEVICE_DO_NOT_ASK else SPARK_SINGLE_DEVICE_DO_NOT_ASK
+          true ->
+            if (isMultiDevice) BLAZE_MULTI_DEVICE_DO_NOT_ASK else BLAZE_SINGLE_DEVICE_DO_NOT_ASK
+        }
+
+      private fun getDialogTitleAndMessage(billingStatus: Boolean?): Pair<String, String> {
+        val devices = devices.value.filterIsInstance<DirectAccessDeviceHandle>()
+
+        return if (devices.isEmpty()) {
+          when (billingStatus) {
+            null ->
+              Pair(
+                "Connect to ${properties.title}",
+                "Devices are reserved for 30 minutes. Unused minutes will be returned. If your project is a Blaze plan you may incur charges.",
+              )
+            false ->
+              Pair(
+                "Connect to ${properties.title}",
+                "Devices are reserved for 30 minutes and count toward your Spark Plan free minutes. When you end your session unused time is refunded.",
+              )
+            true ->
+              Pair(
+                "Connect to ${properties.title}",
+                "You are currently using a Firebase project on the Blaze plan. This session may incur billed usage.",
+              )
+          }
+        } else {
+          when (billingStatus) {
+            null ->
+              Pair(
+                "Reserving Multiple Streaming Devices",
+                "Devices are reserved for 30 minutes. Unused minutes will be returned. If your project is a Blaze plan you may incur charges.",
+              )
+            false ->
+              Pair(
+                "Reserving Multiple Streaming Devices",
+                "You have another streaming device reserved. The new device will be reserved and count towards your Spark Plan free minutes.",
+              )
+            true ->
+              Pair(
+                "Reserving Multiple Streaming Devices",
+                "You have another device streaming session. You are currently using a Firebase project on the Blaze plan. This session may incur billed usage.",
+              )
+          }
+        }
+      }
+
       private suspend fun confirmWaitingTime() {
         deviceInfo.deviceAvailabilityEstimateSeconds
           ?.let { seconds -> waitTimeText(seconds, "minutes") }
diff --git a/directaccess/testSrc/com/google/gct/directaccess/TestUtils.kt b/directaccess/testSrc/com/google/gct/directaccess/TestUtils.kt
index b3088d2..c913679 100644
--- a/directaccess/testSrc/com/google/gct/directaccess/TestUtils.kt
+++ b/directaccess/testSrc/com/google/gct/directaccess/TestUtils.kt
@@ -98,7 +98,7 @@ object TestUtils {
         50,
         100,
         150,
-        30,
+        null,
       ),
     )
   }
diff --git a/directaccess/testSrc/com/google/gct/directaccess/analytics/DirectAccessUsageTrackerTest.kt b/directaccess/testSrc/com/google/gct/directaccess/analytics/DirectAccessUsageTrackerTest.kt
index 7626f75..e5ca152 100644
--- a/directaccess/testSrc/com/google/gct/directaccess/analytics/DirectAccessUsageTrackerTest.kt
+++ b/directaccess/testSrc/com/google/gct/directaccess/analytics/DirectAccessUsageTrackerTest.kt
@@ -49,6 +49,7 @@ import com.google.gct.directaccess.provisioner.DirectAccessDeviceProvisionerPlug
 import com.google.gct.directaccess.provisioner.DirectAccessDeviceTemplate
 import com.google.gct.directaccess.provisioner.PLUGIN_ID
 import com.google.gct.directaccess.rule.CleanUpNotificationRule
+import com.google.gct.directaccess.rule.PropertiesComponentRule
 import com.google.gct.login2.GoogleLoginService
 import com.google.gct.login2.LoginUsersRule
 import com.google.services.firebase.directaccess.client.DirectAccessConnection.ConnectionState
@@ -108,6 +109,7 @@ class DirectAccessUsageTrackerTest {
   private val grpcConnectionRule = GrpcConnectionRule(listOf(service))
   private val loginUsersRule = LoginUsersRule()
   private val cleanUpNotificationRule = CleanUpNotificationRule(projectRule)
+  private val propertiesComponentRule = PropertiesComponentRule(projectRule)
 
   @get:Rule
   val ruleChain: RuleChain =
@@ -116,6 +118,7 @@ class DirectAccessUsageTrackerTest {
       .around(grpcConnectionRule)
       .around(loginUsersRule)
       .around(cleanUpNotificationRule)
+      .around(propertiesComponentRule)
 
   private val session = FakeAdbSession()
   private lateinit var plugin: DirectAccessDeviceProvisionerPlugin
diff --git a/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTest.kt b/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTest.kt
index 6d6e3a0..82ed0ed 100644
--- a/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTest.kt
+++ b/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTest.kt
@@ -62,6 +62,7 @@ import com.google.gct.directaccess.TestUtils.showAllTemplates
 import com.google.gct.directaccess.analytics.DirectAccessUsageTracker
 import com.google.gct.directaccess.rule.CleanUpNotificationRule
 import com.google.gct.directaccess.rule.FakeToolWindowRule
+import com.google.gct.directaccess.rule.PropertiesComponentRule
 import com.google.gct.directaccess.ui.SelectDeviceDialog
 import com.google.gct.login.CredentialedUser
 import com.google.gct.login.GoogleLogin
@@ -132,6 +133,7 @@ class DirectAccessDeviceProvisionerTest {
   private val loginStateRule = LoginStateRule(LoginStatus.LoggedIn("test@gmail.com"))
   private val fakeToolWindowRule = FakeToolWindowRule(projectRule)
   private val cleanUpNotificationRule = CleanUpNotificationRule(projectRule)
+  private val propertiesComponentRule = PropertiesComponentRule(projectRule)
 
   @get:Rule
   val ruleChain: RuleChain =
@@ -141,6 +143,7 @@ class DirectAccessDeviceProvisionerTest {
       .around(loginStateRule)
       .around(fakeToolWindowRule)
       .around(cleanUpNotificationRule)
+      .around(propertiesComponentRule)
 
   private val session = FakeAdbSession()
   private lateinit var plugin: DirectAccessDeviceProvisionerPlugin
diff --git a/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTestWithLogin2.kt b/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTestWithLogin2.kt
index b5abb6a..41cf889 100644
--- a/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTestWithLogin2.kt
+++ b/directaccess/testSrc/com/google/gct/directaccess/provisioner/DirectAccessDeviceProvisionerTestWithLogin2.kt
@@ -61,8 +61,10 @@ import com.google.gct.directaccess.TestUtils.refreshReservations
 import com.google.gct.directaccess.TestUtils.reservation
 import com.google.gct.directaccess.TestUtils.showAllTemplates
 import com.google.gct.directaccess.analytics.DirectAccessUsageTracker
+import com.google.gct.directaccess.directAccessCloudProjectManager
 import com.google.gct.directaccess.rule.CleanUpNotificationRule
 import com.google.gct.directaccess.rule.FakeToolWindowRule
+import com.google.gct.directaccess.rule.PropertiesComponentRule
 import com.google.gct.directaccess.ui.SelectDeviceDialog
 import com.google.gct.login2.GoogleLoginService
 import com.google.gct.login2.LoginUsersRule
@@ -77,6 +79,7 @@ import com.google.services.firebase.directaccess.client.isActive
 import com.google.services.firebase.directaccess.client.waitUntilActive
 import com.google.wireless.android.sdk.stats.DeviceInfo
 import com.intellij.icons.AllIcons
+import com.intellij.ide.util.PropertiesComponent
 import com.intellij.notification.Notification
 import com.intellij.notification.NotificationAction
 import com.intellij.notification.NotificationDisplayType
@@ -99,6 +102,7 @@ import icons.StudioIcons
 import icons.StudioIcons.DeviceExplorer.FIREBASE_DEVICE_PHONE
 import icons.StudioIcons.DeviceExplorer.FIREBASE_DEVICE_WEAR
 import java.time.Duration
+import java.util.concurrent.CountDownLatch
 import java.util.concurrent.TimeUnit
 import javax.swing.Icon
 import javax.swing.JLabel
@@ -131,6 +135,7 @@ class DirectAccessDeviceProvisionerTestWithLogin2 {
   private val loginUsersRule = LoginUsersRule()
   private val fakeToolWindowRule = FakeToolWindowRule(projectRule)
   private val cleanUpNotificationRule = CleanUpNotificationRule(projectRule)
+  private val propertiesComponentRule = PropertiesComponentRule(projectRule)
 
   @get:Rule
   val ruleChain: RuleChain =
@@ -140,6 +145,7 @@ class DirectAccessDeviceProvisionerTestWithLogin2 {
       .around(loginUsersRule)
       .around(fakeToolWindowRule)
       .around(cleanUpNotificationRule)
+      .around(propertiesComponentRule)
 
   private val session = FakeAdbSession()
   private lateinit var plugin: DirectAccessDeviceProvisionerPlugin
@@ -1243,6 +1249,97 @@ class DirectAccessDeviceProvisionerTestWithLogin2 {
     assertThat(service<GoogleLoginService>().isLoggedIn()).isTrue()
   }
 
+  @Test
+  fun testUnknownUsagePromptBeforeReservingDevice() =
+    testUsagePromptBeforeReservingDevice(
+      null,
+      "Devices are reserved for 30 minutes. Unused minutes will be returned. If your project is a Blaze plan you may incur charges.",
+      "Devices are reserved for 30 minutes. Unused minutes will be returned. If your project is a Blaze plan you may incur charges.",
+      UNKNOWN_DEVICE_DO_NOT_ASK,
+      UNKNOWN_DEVICE_DO_NOT_ASK,
+    )
+
+  @Test
+  fun testSparkUsagePromptBeforeReservingDevice() =
+    testUsagePromptBeforeReservingDevice(
+      false,
+      "Devices are reserved for 30 minutes and count toward your Spark Plan free minutes. When you end your session unused time is refunded.",
+      "You have another streaming device reserved. The new device will be reserved and count towards your Spark Plan free minutes.",
+      SPARK_SINGLE_DEVICE_DO_NOT_ASK,
+      SPARK_MULTI_DEVICE_DO_NOT_ASK,
+    )
+
+  @Test
+  fun testBlazeUsagePromptBeforeReservingDevice() =
+    testUsagePromptBeforeReservingDevice(
+      true,
+      "You are currently using a Firebase project on the Blaze plan. This session may incur billed usage.",
+      "You have another device streaming session. You are currently using a Firebase project on the Blaze plan. This session may incur billed usage.",
+      BLAZE_SINGLE_DEVICE_DO_NOT_ASK,
+      BLAZE_MULTI_DEVICE_DO_NOT_ASK,
+    )
+
+  private fun testUsagePromptBeforeReservingDevice(
+    billing: Boolean?,
+    singleMessage: String,
+    multiMessage: String,
+    singleKey: String,
+    multiKey: String,
+  ) =
+    try {
+      StudioFlags.DIRECT_ACCESS_MONTHLY_QUOTA.override(true)
+      runBlockingWithTimeout {
+        val countDownLatch = CountDownLatch(2)
+        // Unset values set by PropertiesComponentRule
+        PropertiesComponent.getInstance(projectRule.project).unsetValue(singleKey)
+        PropertiesComponent.getInstance(projectRule.project).unsetValue(multiKey)
+        val cloudProjectManager = projectRule.project.directAccessCloudProjectManager!!
+        val billingEnabledFlow = RefreshableStateFlow(scope, TimeUnit.HOURS.toMillis(1)) { billing }
+        doAnswer { billingEnabledFlow }.whenever(cloudProjectManager).isBillingEnabledFlow
+
+        TestDialogManager.setTestDialog { message ->
+          assertThat(message).isEqualTo(singleMessage)
+          countDownLatch.countDown()
+          Messages.YES
+        }
+        // Reserve a device
+        plugin.templates.value[0].activationAction.activate()
+        yieldUntil { plugin.devices.value.size == 1 }
+
+        TestDialogManager.setTestDialog { message ->
+          assertThat(message).isEqualTo(multiMessage)
+          countDownLatch.countDown()
+          Messages.YES
+        }
+        // Reserve another device for multi-device prompt
+        plugin.templates.value[4].activationAction.activate()
+        yieldUntil { plugin.devices.value.size == 2 }
+
+        // Check the "do not ask again" box
+        PropertiesComponent.getInstance(projectRule.project).setValue(singleKey, true)
+        PropertiesComponent.getInstance(projectRule.project).setValue(multiKey, true)
+
+        // Setup prompt to fail test if invoked
+        TestDialogManager.setTestDialog {
+          fail("Should not prompt")
+          Messages.YES
+        }
+
+        plugin.devices.value.forEach { it.deactivationAction?.deactivate() }
+        yieldUntil { plugin.devices.value.isEmpty() }
+
+        // Reserve devices again
+        plugin.templates.value[0].activationAction.activate()
+        yieldUntil { plugin.devices.value.size == 1 }
+
+        plugin.templates.value[4].activationAction.activate()
+        yieldUntil { plugin.devices.value.size == 2 }
+        assertThat(countDownLatch.count).isEqualTo(0)
+      }
+    } finally {
+      StudioFlags.DIRECT_ACCESS_MONTHLY_QUOTA.clearOverride()
+    }
+
   private suspend fun testCorrectIcon(template: DirectAccessDeviceTemplate, icon: Icon) {
     val handle = template.activationAction.activate() as DirectAccessDeviceHandle
     session.hostServices.connect(handle.connection.deviceAddress()!!)
diff --git a/directaccess/testSrc/com/google/gct/directaccess/rule/PropertiesComponentRule.kt b/directaccess/testSrc/com/google/gct/directaccess/rule/PropertiesComponentRule.kt
new file mode 100644
index 0000000..e9ae30a
--- /dev/null
+++ b/directaccess/testSrc/com/google/gct/directaccess/rule/PropertiesComponentRule.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2024 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.google.gct.directaccess.rule
+
+import com.google.gct.directaccess.provisioner.BLAZE_MULTI_DEVICE_DO_NOT_ASK
+import com.google.gct.directaccess.provisioner.BLAZE_SINGLE_DEVICE_DO_NOT_ASK
+import com.google.gct.directaccess.provisioner.SPARK_MULTI_DEVICE_DO_NOT_ASK
+import com.google.gct.directaccess.provisioner.SPARK_SINGLE_DEVICE_DO_NOT_ASK
+import com.google.gct.directaccess.provisioner.UNKNOWN_DEVICE_DO_NOT_ASK
+import com.intellij.ide.util.PropertiesComponent
+import com.intellij.testFramework.ProjectRule
+import org.junit.rules.ExternalResource
+
+private val keyList =
+  listOf(
+    UNKNOWN_DEVICE_DO_NOT_ASK,
+    SPARK_SINGLE_DEVICE_DO_NOT_ASK,
+    SPARK_MULTI_DEVICE_DO_NOT_ASK,
+    BLAZE_SINGLE_DEVICE_DO_NOT_ASK,
+    BLAZE_MULTI_DEVICE_DO_NOT_ASK,
+  )
+
+class PropertiesComponentRule(private val projectRule: ProjectRule) : ExternalResource() {
+  override fun before() {
+    keyList.forEach { PropertiesComponent.getInstance(projectRule.project).setValue(it, true) }
+  }
+
+  override fun after() {
+    keyList.forEach { PropertiesComponent.getInstance(projectRule.project).unsetValue(it) }
+  }
+}