Jetpack REST connection: Finalize

WordPress/src/main/java/org/wordpress/android/ui/jetpackrestconnection/JetpackConnector.kt

@@ -6,7 +6,7 @@ import uniffi.wp_api.WpComSiteId
 import javax.inject.Inject
 
 class JetpackConnector @Inject constructor(
-    private val jetpackConnectionHelper: JetpackConnectionHelper
+    private val jetpackConnectionHelper: JetpackConnectionHelper,
 ) {
     /**
      * Connects the Jetpack site to WordPress.com and returns the site ID

WordPress/src/main/java/org/wordpress/android/ui/jetpackrestconnection/JetpackRestConnectionActivity.kt

@@ -2,6 +2,7 @@ package org.wordpress.android.ui.jetpackrestconnection
 
 import android.content.Context
 import android.content.Intent
+import android.os.Build
 import android.os.Bundle
 import androidx.activity.viewModels
 import androidx.compose.runtime.collectAsState
@@ -15,6 +16,7 @@ import org.wordpress.android.fluxc.store.AccountStore
 import org.wordpress.android.ui.ActivityLauncher
 import org.wordpress.android.ui.ActivityNavigator
 import org.wordpress.android.ui.RequestCodes
+import org.wordpress.android.ui.jetpackrestconnection.JetpackRestConnectionViewModel.Companion.DEFAULT_CONNECTION_SOURCE
 import org.wordpress.android.ui.main.BaseAppCompatActivity
 import org.wordpress.android.util.extensions.setContent
 import javax.inject.Inject
@@ -31,12 +33,28 @@ class JetpackRestConnectionActivity : BaseAppCompatActivity() {
 
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
+
+        // Get the connection source from the intent and set it in the ViewModel
+        if (savedInstanceState == null) {
+            val source = if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.TIRAMISU) {
+                intent.getSerializableExtra(
+                    KEY_CONNECTION_SOURCE,
+                    JetpackRestConnectionViewModel.ConnectionSource::class.java
+                )
+            } else {
+                @Suppress("DEPRECATION")
+                intent.getSerializableExtra(KEY_CONNECTION_SOURCE) as? JetpackRestConnectionViewModel.ConnectionSource
+            } ?: DEFAULT_CONNECTION_SOURCE
+            viewModel.setConnectionSource(source)
+        }
+
         setContent {
             JetpackRestConnectionScreen(
                 currentStep = viewModel.currentStep.collectAsState(),
                 stepStates = viewModel.stepStates.collectAsState(),
                 buttonType = viewModel.buttonType.collectAsState(),
                 onStartClick = viewModel::onStartClick,
+                onDoneClick = viewModel::onDoneClick,
                 onCloseClick = viewModel::onCloseClick,
                 onRetryClick = viewModel::onRetryClick
             )
@@ -48,6 +66,10 @@ class JetpackRestConnectionActivity : BaseAppCompatActivity() {
                     JetpackRestConnectionViewModel.UiEvent.StartWPComLogin ->
                         startWPComLogin()
 
+                    JetpackRestConnectionViewModel.UiEvent.Done -> {
+                        finish()
+                    }
+
                     JetpackRestConnectionViewModel.UiEvent.Close ->
                         finish()
 
@@ -84,8 +106,17 @@ class JetpackRestConnectionActivity : BaseAppCompatActivity() {
     }
 
     companion object {
-        @JvmStatic
-        fun createIntent(context: Context) =
-            Intent(context, JetpackRestConnectionActivity::class.java)
+        private const val KEY_CONNECTION_SOURCE = "key_connection_source"
+
+        fun startJetpackRestConnectionFlow(
+            context: Context,
+            source: JetpackRestConnectionViewModel.ConnectionSource
+        ) {
+            context.startActivity(
+                Intent(context, JetpackRestConnectionActivity::class.java).apply {
+                    putExtra(KEY_CONNECTION_SOURCE, source)
+                }
+            )
+        }
     }
 }

WordPress/src/main/java/org/wordpress/android/ui/jetpackrestconnection/JetpackRestConnectionScreen.kt

@@ -65,6 +65,7 @@ fun JetpackRestConnectionScreen(
     stepStates: State<Map<ConnectionStep, StepState>>,
     buttonType: State<ButtonType?>,
     onStartClick: () -> Unit = {},
+    onDoneClick: () -> Unit = {},
     onCloseClick: () -> Unit = {},
     onRetryClick: () -> Unit = {}
 ) {
@@ -88,7 +89,7 @@ fun JetpackRestConnectionScreen(
                     JetpackConnectionButton(
                         buttonType = buttonType.value,
                         onStartClick = onStartClick,
-                        onDoneClick = onCloseClick,
+                        onDoneClick = onDoneClick,
                         onRetryClick = onRetryClick
                     )
                 }
@@ -274,6 +275,7 @@ private fun getErrorText(context: Context, errorType: ErrorType): String {
         ErrorType.MissingAccessToken -> R.string.jetpack_rest_connection_error_access_token
         ErrorType.ConnectSiteFailed -> R.string.jetpack_rest_connection_error_connect_site
         ErrorType.InstallJetpackFailed -> R.string.jetpack_rest_connection_error_install_jetpack
+        ErrorType.ActivateStatsFailed -> R.string.jetpack_rest_connection_error_activate_stats
         ErrorType.Timeout -> R.string.jetpack_rest_connection_error_timeout
         ErrorType.Offline -> R.string.jetpack_rest_connection_error_offline
         is ErrorType.Unknown -> R.string.jetpack_rest_connection_error_unknown

WordPress/src/main/java/org/wordpress/android/ui/jetpackrestconnection/JetpackRestConnectionViewModel.kt

@@ -25,10 +25,11 @@ import javax.inject.Named
 class JetpackRestConnectionViewModel @Inject constructor(
     @Named(UI_THREAD) private val mainDispatcher: CoroutineDispatcher,
     @Named(BG_THREAD) private val bgDispatcher: CoroutineDispatcher,
-    private val selectedSiteRepository: SelectedSiteRepository,
+    selectedSiteRepository: SelectedSiteRepository,
     private val accountStore: AccountStore,
     private val jetpackInstaller: JetpackInstaller,
     private val jetpackConnector: JetpackConnector,
+    private val jetpackModuleHelper: JetpackStatsModuleHelper,
     private val appLogWrapper: AppLogWrapper,
 ) : ScopedViewModel(mainDispatcher) {
     private val _currentStep = MutableStateFlow<ConnectionStep?>(null)
@@ -46,6 +47,17 @@ class JetpackRestConnectionViewModel @Inject constructor(
     private var job: Job? = null
     private var isWaitingForWPComLogin = false
 
+    private var connectionSource: ConnectionSource = DEFAULT_CONNECTION_SOURCE
+    private var site: SiteModel = selectedSiteRepository.getSelectedSite() ?: error("No site selected")
+
+    /**
+     * This will be used for analytics tracking
+     */
+    fun setConnectionSource(source: ConnectionSource) {
+        connectionSource = source
+        appLogWrapper.d(AppLog.T.API, "$TAG: Connection source set to: $source")
+    }
+
     private fun startConnectionJob(fromStep: ConnectionStep? = null) {
         val stepInfo = fromStep?.let { " from step: $it" } ?: ""
         appLogWrapper.d(AppLog.T.API, "$TAG: Starting Jetpack connection job$stepInfo")
@@ -145,6 +157,14 @@ class JetpackRestConnectionViewModel @Inject constructor(
         startConnectionJob()
     }
 
+    /**
+     * Connection flow completed successfully and user clicked the Done button
+     */
+    fun onDoneClick() {
+        appLogWrapper.d(AppLog.T.API, "$TAG: Done clicked")
+        setUiEvent(UiEvent.Done)
+    }
+
     /**
      * User clicked the close button, confirm closing if the connection is in progress, otherwise close immediately
      */
@@ -267,13 +287,13 @@ class JetpackRestConnectionViewModel @Inject constructor(
 
             ConnectionStep.Finalize -> {
                 appLogWrapper.d(AppLog.T.API, "$TAG: Finalizing connection")
-                // TODO
+                finalize()
             }
         }
     }
 
     /**
-     * Starts the wp.com login flow if the user isn't logged into wp.com
+     * Step 1: Starts the wp.com login flow if the user isn't logged into wp.com
      */
     private fun loginWpCom() {
         if (accountStore.hasAccessToken()) {
@@ -314,10 +334,10 @@ class JetpackRestConnectionViewModel @Inject constructor(
     }
 
     /**
-     * Installs Jetpack to the current site if not already installed
+     * Step 2: Installs Jetpack to the current site if not already installed
      */
     private suspend fun installJetpack() {
-        val result = jetpackInstaller.installJetpack(getSite())
+        val result = jetpackInstaller.installJetpack(site)
 
         result.fold(
             onSuccess = { status ->
@@ -350,12 +370,15 @@ class JetpackRestConnectionViewModel @Inject constructor(
     }
 
     /**
-     * Connects the current site to Jetpack
+     * Step 3: Connects the current site to Jetpack
      */
     private suspend fun connectSite() {
-        val result = jetpackConnector.connectSite(getSite())
+        val result = jetpackConnector.connectSite(site)
         result.fold(
-            onSuccess = {
+            onSuccess = { wpComSiteId ->
+                // the local site won't have a siteId since it's self-hosted and previously unconnected to Jetpack,
+                // so assign it the siteId retrieved when connecting the site
+                site.siteId = wpComSiteId.toLong()
                 updateStepStatus(
                     step = ConnectionStep.ConnectSite,
                     status = ConnectionStatus.Completed
@@ -372,7 +395,7 @@ class JetpackRestConnectionViewModel @Inject constructor(
     }
 
     /**
-     * Connects the user to the current site to Jetpack
+     * Step 4: Connects the user to the current site to Jetpack
      */
     private suspend fun connectUser() {
         if (!accountStore.hasAccessToken()) {
@@ -384,11 +407,11 @@ class JetpackRestConnectionViewModel @Inject constructor(
             return
         }
         val result = jetpackConnector.connectUser(
-            site = getSite(),
+            site = site,
             accessToken = accountStore.accessToken!!
         )
         result.fold(
-            onSuccess = {
+            onSuccess = { wpComUserId ->
                 updateStepStatus(
                     step = ConnectionStep.ConnectUser,
                     status = ConnectionStatus.Completed
@@ -405,10 +428,23 @@ class JetpackRestConnectionViewModel @Inject constructor(
     }
 
     /**
-     * Gets the current site from the store
+     * Step 5: Finalize the connection by activating the stats module for the site
      */
-    private fun getSite() =
-        selectedSiteRepository.getSelectedSite() ?: error("No site is currently selected in SelectedSiteRepository")
+    private suspend fun finalize() {
+        val result = jetpackModuleHelper.activateStatsModule(site)
+        if (result.isSuccess) {
+            updateStepStatus(
+                step = ConnectionStep.Finalize,
+                status = ConnectionStatus.Completed
+            )
+        } else {
+            updateStepStatus(
+                step = ConnectionStep.Finalize,
+                status = ConnectionStatus.Failed,
+                error = ErrorType.ActivateStatsFailed
+            )
+        }
+    }
 
     sealed class ConnectionStep {
         data object LoginWpCom : ConnectionStep()
@@ -427,6 +463,7 @@ class JetpackRestConnectionViewModel @Inject constructor(
 
     sealed class UiEvent {
         data object StartWPComLogin : UiEvent()
+        data object Done : UiEvent()
         data object Close : UiEvent()
         data object ShowCancelConfirmation : UiEvent()
     }
@@ -438,6 +475,7 @@ class JetpackRestConnectionViewModel @Inject constructor(
         data object ConnectWpComFailed : ErrorType()
         data object ConnectSiteFailed : ErrorType()
         data object ConnectUserFailed : ErrorType()
+        data object ActivateStatsFailed : ErrorType()
         data object MissingAccessToken : ErrorType()
         data object Timeout : ErrorType()
         data object Offline : ErrorType()
@@ -455,11 +493,17 @@ class JetpackRestConnectionViewModel @Inject constructor(
         val errorType: ErrorType? = null,
     )
 
+    enum class ConnectionSource {
+        STATS,
+        NOTIFS
+    }
+
     companion object {
         private const val TAG = "JetpackRestConnectionViewModel"
         private const val LIMIT_VERSION = "14.2"
         private const val STEP_TIMEOUT_MS = 45 * 1000L
         private const val UI_DELAY_MS = 1000L
+        val DEFAULT_CONNECTION_SOURCE = ConnectionSource.STATS
 
         /**
          * Requirements:

WordPress/src/main/java/org/wordpress/android/ui/jetpackrestconnection/JetpackStatsModuleHelper.kt

@@ -0,0 +1,51 @@
+package org.wordpress.android.ui.jetpackrestconnection
+
+import org.wordpress.android.fluxc.model.SiteModel
+import org.wordpress.android.fluxc.store.JetpackStore
+import org.wordpress.android.fluxc.store.JetpackStore.ActivateStatsModulePayload
+import org.wordpress.android.fluxc.store.SiteStore
+import javax.inject.Inject
+
+class JetpackStatsModuleHelper @Inject constructor(
+    private val jetpackStore: JetpackStore,
+    private val siteStore: SiteStore
+) {
+    suspend fun activateStatsModule(site: SiteModel): Result<Unit> {
+        if (isStatsModuleActivated(site)) {
+            return Result.success(Unit)
+        }
+
+        jetpackStore.activateStatsModule(ActivateStatsModulePayload(site))
+
+        // ignore the activation response as it's unreliable, instead just check if it's activated now
+        return verifyStatsModuleActivation(site)
+    }
+
+    /**
+     * Fetch Jetpack sites filtering for this site and check if the module is now activated. Note that fetching
+     * the single site fails because of changes to it during the connection flow.
+     */
+    private suspend fun verifyStatsModuleActivation(
+        site: SiteModel,
+    ): Result<Unit> {
+        val payload = SiteStore.FetchSitesPayload(
+            filters = listOf(SiteStore.SiteFilter.JETPACK)
+        )
+        return siteStore.fetchSites(payload).updatedSites.firstOrNull {
+            it.siteId == site.siteId
+        }?.let { updatedSite ->
+            if (isStatsModuleActivated(updatedSite)) {
+                Result.success(Unit)
+            } else {
+                Result.failure(Exception("Stats module not activated"))
+            }
+        } ?: Result.failure(Exception("Site not found"))
+    }
+
+    private fun isStatsModuleActivated(site: SiteModel) =
+        site.isActiveModuleEnabled(STATS_MODULE_NAME)
+
+    companion object {
+        private const val STATS_MODULE_NAME = "stats"
+    }
+}