Make it possible to run feature specific tasks on login completedmirror-goog-studio-main

...if the user logs in through the settings page or the combined
sign-in ui.

Bug: 334152239
Test: added
Change-Id: I342e87b0afe2db86fd405d334d8a037b70ecc74e

10 files changed, 110 insertions, 28 deletions

diff --git a/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt b/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt
index 3035d93..28fb289 100644
--- a/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt
+++ b/google-login-plugin/src/com/google/gct/login2/GoogleLoginService.kt
@@ -82,6 +82,7 @@ interface GoogleLoginService {
   suspend fun logIn(
     features: Set<LoginFeature> = emptySet(),
     preferredUser: PreferredUser = PreferredUser.ActiveUser,
+    loginCompletedCallback: LoginLogoutCompletedCallback? = null,
     parentComponent: JComponent? = null,
   ): Boolean
 
@@ -238,9 +239,7 @@ internal class GoogleLoginServiceImpl(
   ) {
     if (useOldVersion) GoogleLogin.instance.logIn(null) { loginCompletedCallback?.onCompleted() }
     else {
-      scope
-        .launch { logIn(features, preferredUser, parentComponent) }
-        .invokeOnCompletion { loginCompletedCallback?.onCompleted() }
+      scope.launch { logIn(features, preferredUser, loginCompletedCallback, parentComponent) }
     }
   }
 
@@ -281,6 +280,7 @@ internal class GoogleLoginServiceImpl(
   override suspend fun logIn(
     features: Set<LoginFeature>,
     preferredUser: PreferredUser,
+    loginCompletedCallback: LoginLogoutCompletedCallback?,
     parentComponent: JComponent?,
   ): Boolean {
     if (useOldVersion) {
@@ -327,6 +327,7 @@ internal class GoogleLoginServiceImpl(
               userManager.addOrUpdateUser(checkNotNull(state.createCredentialedUser()))
 
               maybeSetActiveUser()
+              loginCompletedCallback?.onCompleted()
 
               cont.resume(true)
             } else {
diff --git a/google-login-plugin/src/com/google/gct/login2/LoginFeature.kt b/google-login-plugin/src/com/google/gct/login2/LoginFeature.kt
index ac5cf05..df25568 100644
--- a/google-login-plugin/src/com/google/gct/login2/LoginFeature.kt
+++ b/google-login-plugin/src/com/google/gct/login2/LoginFeature.kt
@@ -47,6 +47,9 @@ interface LoginFeature {
   val settingsAction: AnAction?
   val oAuthScopes: Collection<String>
   val onboardingWizardEntry: OnboardingWizardEntry?
+  val onLoginCompleted: LoginLogoutCompletedCallback?
+    get() = null
+
   val isInternal: Boolean
     get() = false
 
diff --git a/google-login-plugin/src/com/google/gct/login2/LoginLogoutCompletedCallback.java b/google-login-plugin/src/com/google/gct/login2/LoginLogoutCompletedCallback.kt
index 17362d4..a0ef92f 100644
--- a/google-login-plugin/src/com/google/gct/login2/LoginLogoutCompletedCallback.java
+++ b/google-login-plugin/src/com/google/gct/login2/LoginLogoutCompletedCallback.kt
@@ -13,14 +13,21 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.google.gct.login2;
+package com.google.gct.login2
 
-/**
- * Callback for when a login or a logout is completed.
- */
-public interface LoginLogoutCompletedCallback {
-  /**
-   * Called when log in is complete.
-   */
-  void onCompleted();
+/** Callback for when a login or a logout is completed. */
+fun interface LoginLogoutCompletedCallback {
+  /** Called when log in is complete. */
+  fun onCompleted()
+
+  operator fun plus(other: LoginLogoutCompletedCallback?): LoginLogoutCompletedCallback {
+    return LoginLogoutCompletedCallback {
+      onCompleted()
+      other?.onCompleted()
+    }
+  }
+
+  companion object {
+    val EMPTY = LoginLogoutCompletedCallback {}
+  }
 }
diff --git a/google-login-plugin/src/com/google/gct/login2/settings/GoogleLoginApplicationSettingsConfigurableUi.kt b/google-login-plugin/src/com/google/gct/login2/settings/GoogleLoginApplicationSettingsConfigurableUi.kt
index 5ec5e39..7f71d19 100644
--- a/google-login-plugin/src/com/google/gct/login2/settings/GoogleLoginApplicationSettingsConfigurableUi.kt
+++ b/google-login-plugin/src/com/google/gct/login2/settings/GoogleLoginApplicationSettingsConfigurableUi.kt
@@ -250,7 +250,9 @@ internal class GoogleLoginApplicationSettingsConfigurableUi : ConfigurableUi<Uni
       JButton().apply {
         addActionListener {
           feature.logInAsync(
-            selectedUserFlow.value?.email?.let { PreferredUser.User(it) } ?: PreferredUser.None,
+            preferredUser =
+              selectedUserFlow.value?.email?.let { PreferredUser.User(it) } ?: PreferredUser.None,
+            loginCompletedCallback = feature.onLoginCompleted,
             parentComponent = this@apply,
           )
         }
diff --git a/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginAction.kt b/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginAction.kt
index fccb6e5..b489d17 100644
--- a/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginAction.kt
+++ b/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginAction.kt
@@ -45,7 +45,7 @@ class GoogleLoginAction : AnAction(), RightAlignedToolbarAction {
   override fun actionPerformed(event: AnActionEvent) {
     lateinit var popup: JBPopup
     val closePopup = { popup.closeOk(null) }
-    val usersPanel = GoogleLoginUsersPanel(closePopup)
+    val usersPanel = GoogleLoginUsersPanel(event.project, closePopup)
     val popupBuilder = JBPopupFactory.getInstance().createComponentPopupBuilder(usersPanel, null)
     popup = popupBuilder.setCancelOnWindowDeactivation(true).createPopup()
     val inputEvent = event.inputEvent
diff --git a/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginUsersPanel.kt b/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginUsersPanel.kt
index 5452565..4e664aa 100644
--- a/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginUsersPanel.kt
+++ b/google-login-plugin/src/com/google/gct/login2/ui/GoogleLoginUsersPanel.kt
@@ -25,6 +25,7 @@ import com.google.gct.login2.ui.onboarding.GoogleSignInWizard
 import com.intellij.ide.ui.laf.darcula.ui.DarculaButtonUI.DEFAULT_STYLE_KEY
 import com.intellij.openapi.components.service
 import com.intellij.openapi.options.ShowSettingsUtil
+import com.intellij.openapi.project.Project
 import com.intellij.openapi.ui.VerticalFlowLayout
 import com.intellij.ui.NewUI
 import com.intellij.ui.RelativeFont
@@ -63,7 +64,8 @@ import javax.swing.SwingConstants
  * The Google Login Panel that displays the currently logged-in users and buttons to add a new user
  * and sign out a logged-in user.
  */
-class GoogleLoginUsersPanel(private val closePopup: () -> Unit) : JPanel() {
+class GoogleLoginUsersPanel(private val project: Project?, private val closePopup: () -> Unit) :
+  JPanel() {
   private fun panel(layout: LayoutManager) =
     JPanel(layout).apply {
       isOpaque = true
@@ -167,7 +169,7 @@ class GoogleLoginUsersPanel(private val closePopup: () -> Unit) : JPanel() {
             override fun mouseClicked(e: MouseEvent) {
               closePopup()
               ShowSettingsUtil.getInstance()
-                .showSettingsDialog(null, GOOGLE_LOGIN_CONFIGURABLE_NAME)
+                .showSettingsDialog(project, GOOGLE_LOGIN_CONFIGURABLE_NAME)
             }
 
             override fun mouseEntered(e: MouseEvent?) {
@@ -209,7 +211,7 @@ class GoogleLoginUsersPanel(private val closePopup: () -> Unit) : JPanel() {
             ActionLink("Enable integrations...") {
                 closePopup()
                 ShowSettingsUtil.getInstance()
-                  .showSettingsDialog(null, GOOGLE_LOGIN_CONFIGURABLE_NAME)
+                  .showSettingsDialog(project, GOOGLE_LOGIN_CONFIGURABLE_NAME)
               }
               .apply { border = JBUI.Borders.emptyLeft(19) }
           )
@@ -267,7 +269,7 @@ class GoogleLoginUsersPanel(private val closePopup: () -> Unit) : JPanel() {
                   GoogleSignInWizard().show()
                 } else {
                   ShowSettingsUtil.getInstance()
-                    .showSettingsDialog(null, GOOGLE_LOGIN_CONFIGURABLE_NAME)
+                    .showSettingsDialog(project, GOOGLE_LOGIN_CONFIGURABLE_NAME)
                 }
               }
             }
diff --git a/google-login-plugin/src/com/google/gct/login2/ui/onboarding/CloudServicesChooserStep.kt b/google-login-plugin/src/com/google/gct/login2/ui/onboarding/CloudServicesChooserStep.kt
index 0084af3..fe54d8b 100644
--- a/google-login-plugin/src/com/google/gct/login2/ui/onboarding/CloudServicesChooserStep.kt
+++ b/google-login-plugin/src/com/google/gct/login2/ui/onboarding/CloudServicesChooserStep.kt
@@ -19,6 +19,7 @@ import com.android.tools.idea.wizard.model.ModelWizard
 import com.android.tools.idea.wizard.model.ModelWizardStep
 import com.google.gct.login2.GoogleLoginService
 import com.google.gct.login2.LoginFeature
+import com.google.gct.login2.LoginLogoutCompletedCallback
 import com.google.gct.login2.PreferredUser
 import com.intellij.openapi.ui.DialogPanel
 import com.intellij.ui.JBColor
@@ -47,6 +48,10 @@ class CloudServicesChooserStep(model: GoogleSignInModel) :
       GoogleLoginService.instance.logIn(
         features = model.requiredIntegrations,
         preferredUser = PreferredUser.None,
+        loginCompletedCallback =
+          model.requiredIntegrations
+            .map { it.onLoginCompleted }
+            .fold(LoginLogoutCompletedCallback.EMPTY) { acc, callback -> acc + callback },
         parentComponent = dialogPanel,
       )
     }
diff --git a/google-login-plugin/testSrc/com/google/gct/login2/GoogleLoginServiceTest.kt b/google-login-plugin/testSrc/com/google/gct/login2/GoogleLoginServiceTest.kt
index 4951c80..d9f40a3 100644
--- a/google-login-plugin/testSrc/com/google/gct/login2/GoogleLoginServiceTest.kt
+++ b/google-login-plugin/testSrc/com/google/gct/login2/GoogleLoginServiceTest.kt
@@ -49,6 +49,7 @@ import com.intellij.util.application
 import java.io.IOException
 import java.util.concurrent.CountDownLatch
 import java.util.concurrent.Semaphore
+import java.util.concurrent.TimeUnit
 import kotlin.time.Duration.Companion.seconds
 import kotlinx.coroutines.runBlocking
 import org.junit.After
@@ -192,6 +193,33 @@ class GoogleLoginServiceTest {
   }
 
   @Test
+  fun `can have custom task performed on login completed, log in synchronously`() {
+    val countdownLatch = CountDownLatch(1)
+    val task = LoginLogoutCompletedCallback { countdownLatch.countDown() }
+
+    val loginService = initGoogleLoginService()
+    runBlocking {
+      loginService.logIn(features = setOf(loginFeatureRule.FEATURE1), loginCompletedCallback = task)
+    }
+    assertThat(settings.activeUser).isEqualTo(TEST_EMAIL)
+    assertThat(countdownLatch.await(1, TimeUnit.SECONDS)).isTrue()
+  }
+
+  @Test
+  fun `can have custom task performed on login completed, log in asynchronously`() {
+    val countdownLatch = CountDownLatch(1)
+    val task = LoginLogoutCompletedCallback { countdownLatch.countDown() }
+
+    val loginService = initGoogleLoginService()
+    loginService.logInAsync(
+      features = setOf(loginFeatureRule.FEATURE1),
+      loginCompletedCallback = task,
+    )
+    assertThat(countdownLatch.await(1, TimeUnit.SECONDS)).isTrue()
+    assertThat(settings.activeUser).isEqualTo(TEST_EMAIL)
+  }
+
+  @Test
   fun `track failed to logout single user`() {
     val errorOAuthServer =
       object : FakeOAuthServer() {
@@ -948,7 +976,7 @@ class GoogleLoginServiceTestWithOldVersion {
     newService.setActiveUser("foo")
     verify(oldLoginService).setActiveUser("foo")
 
-    val callback = LoginLogoutCompletedCallback {}
+    val callback = LoginLogoutCompletedCallback.EMPTY
     newService.logInAsync(
       setOf(loginFeatureRule.FEATURE1),
       PreferredUser.ActiveUser,
diff --git a/google-login-plugin/testSrc/com/google/gct/login2/LoginLogoutCompletedCallbackTest.kt b/google-login-plugin/testSrc/com/google/gct/login2/LoginLogoutCompletedCallbackTest.kt
new file mode 100644
index 0000000..60459b0
--- /dev/null
+++ b/google-login-plugin/testSrc/com/google/gct/login2/LoginLogoutCompletedCallbackTest.kt
@@ -0,0 +1,33 @@
+/*
+ * 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.login2
+
+import com.google.common.truth.Truth.assertThat
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import org.junit.Test
+
+class LoginLogoutCompletedCallbackTest {
+  @Test
+  fun `plus callbacks`() {
+    val countDownLatch = CountDownLatch(2)
+    val callback1 = LoginLogoutCompletedCallback { countDownLatch.countDown() }
+    val callback2 = LoginLogoutCompletedCallback { countDownLatch.countDown() }
+
+    (callback1 + callback2).onCompleted()
+    assertThat(countDownLatch.await(1, TimeUnit.SECONDS)).isTrue()
+  }
+}
diff --git a/google-login-plugin/testSrc/com/google/gct/login2/ui/GoogleLoginUsersPanelTest.kt b/google-login-plugin/testSrc/com/google/gct/login2/ui/GoogleLoginUsersPanelTest.kt
index e90a275..9278716 100644
--- a/google-login-plugin/testSrc/com/google/gct/login2/ui/GoogleLoginUsersPanelTest.kt
+++ b/google-login-plugin/testSrc/com/google/gct/login2/ui/GoogleLoginUsersPanelTest.kt
@@ -33,9 +33,9 @@ import com.intellij.ide.ui.search.SearchUtil
 import com.intellij.openapi.application.ApplicationManager
 import com.intellij.openapi.options.ShowSettingsUtil
 import com.intellij.openapi.options.ex.ConfigurableExtensionPointUtil
-import com.intellij.testFramework.ApplicationRule
 import com.intellij.testFramework.DisposableRule
 import com.intellij.testFramework.EdtRule
+import com.intellij.testFramework.ProjectRule
 import com.intellij.testFramework.RunsInEdt
 import com.intellij.testFramework.replaceService
 import com.intellij.ui.components.ActionLink
@@ -70,6 +70,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
   private val TEST_DATA_ROOT =
     TestUtils.resolveWorkspacePath("tools/studio/google/cloud/tools/google-login-plugin/testData")
 
+  private val projectRule = ProjectRule()
   private val disposableRule = DisposableRule()
   private val loginUsersRule = LoginUsersRule()
   private val loginFeatureRule = LoginFeatureRule()
@@ -79,7 +80,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
     RuleChain.outerRule(FlagRule(StudioFlags.ENABLE_SETTINGS_ACCOUNT_UI, true))
       .around(FlagRule(StudioFlags.ENABLE_COMBINED_LOGIN_UI, combinedLoginFlagOnOff))
       .around(disposableRule)
-      .around(ApplicationRule())
+      .around(projectRule)
       .around(EdtRule())
       .around(loginFeatureRule)
       .around(loginUsersRule)!!
@@ -90,7 +91,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
       .replaceService(ShowSettingsUtil::class.java, settingsUtil, disposableRule.disposable)
 
     var closeInvoked = false
-    val panel = GoogleLoginUsersPanel { closeInvoked = true }
+    val panel = GoogleLoginUsersPanel(projectRule.project) { closeInvoked = true }
 
     val invokeAction = testBody(panel)
 
@@ -99,7 +100,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
     assertThat(closeInvoked).isTrue()
 
     // validate that we're trying to show the correct configurable
-    Mockito.verify(settingsUtil).showSettingsDialog(null, "Google Accounts")
+    Mockito.verify(settingsUtil).showSettingsDialog(projectRule.project, "Google Accounts")
 
     // validate that the accounts configurable actually exists
     val groups = ConfigurableExtensionPointUtil.getConfigurableGroup(null, true)
@@ -109,7 +110,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
   private fun validateCombinedLoginUiOpenedTest(testBody: (GoogleLoginUsersPanel) -> (() -> Unit)) {
     enableHeadlessDialogs(disposableRule.disposable)
     var closeInvoked = false
-    val panel = GoogleLoginUsersPanel { closeInvoked = true }
+    val panel = GoogleLoginUsersPanel(projectRule.project) { closeInvoked = true }
 
     val invokeAction = testBody(panel)
 
@@ -164,7 +165,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
   fun testSignOut() {
     loginUsersRule.setActiveUser("foo@example.com")
     loginUsersRule.setActiveUser("bar@example.com")
-    val panel = GoogleLoginUsersPanel {}
+    val panel = GoogleLoginUsersPanel(projectRule.project) {}
     val link = panel.flatten().find { (it as? JLabel)?.text == "Sign Out All" } as JLabel
     link.parent.mouseListeners.forEach { it.mouseClicked(mock()) }
     UIUtil.dispatchAllInvocationEvents()
@@ -182,7 +183,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
     )
     loginUsersRule.setActiveUser("mainUser@example.com")
 
-    val panel = GoogleLoginUsersPanel {}
+    val panel = GoogleLoginUsersPanel(projectRule.project) {}
     val list = panel.flatten().filterIsInstance<JBList<*>>().first()
     val foo = loginUsersRule.getUser("foo@example.com")
     assertThat((list.model as DefaultListModel<*>).elements().toList())
@@ -223,7 +224,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
       picture = createUserImage(),
       features = listOf(loginFeatureRule.FEATURE2),
     )
-    val panel = GoogleLoginUsersPanel {}
+    val panel = GoogleLoginUsersPanel(projectRule.project) {}
 
     val panelContents = panel.flatten()
     val activeUserPanelContents = panelContents.single { it.name == "ActiveUserPanel" }.flatten()
@@ -245,7 +246,7 @@ class GoogleLoginUsersPanelTest(private val combinedLoginFlagOnOff: Boolean) {
       .isTrue()
 
     loginUsersRule.setActiveUser("u2@example.com", features = listOf())
-    val panel2 = GoogleLoginUsersPanel {}
+    val panel2 = GoogleLoginUsersPanel(projectRule.project) {}
 
     val panel2Contents = panel2.flatten()
     val activeUserPanel2Contents = panel2Contents.single { it.name == "ActiveUserPanel" }.flatten()