Fix data collection crash when UI state isn't Ready

app/src/main/java/org/groundplatform/android/ui/datacollection/DataCollectionFragment.kt

@@ -22,14 +22,14 @@ import android.view.ViewGroup
 import androidx.hilt.navigation.fragment.hiltNavGraphViewModels
 import androidx.navigation.fragment.findNavController
 import dagger.hilt.android.AndroidEntryPoint
+import javax.inject.Inject
 import org.groundplatform.android.R
 import org.groundplatform.android.ui.common.AbstractFragment
 import org.groundplatform.android.ui.common.BackPressListener
 import org.groundplatform.android.ui.common.EphemeralPopups
 import org.groundplatform.android.ui.home.HomeScreenViewModel
 import org.groundplatform.android.util.createComposeView
 import org.groundplatform.android.util.openAppSettings
-import javax.inject.Inject
 
 /** Fragment allowing the user to collect data to complete a task. */
 @AndroidEntryPoint
@@ -71,17 +71,7 @@ class DataCollectionFragment : AbstractFragment(), BackPressListener {
   }
 
   override fun onBack(): Boolean {
-    if (viewModel.uiState.value is DataCollectionUiState.TaskSubmitted) {
-      // Pressing back button after submitting task should navigate back to home screen.
-      navigateBack()
-      return true
-    }
-
-    if (viewModel.isAtFirstTask()) {
-      viewModel.showExitWarning()
-    } else {
-      viewModel.moveToPreviousTask()
-    }
+    viewModel.onBackClicked()
     return true
   }
 

app/src/main/java/org/groundplatform/android/ui/datacollection/DataCollectionViewModel.kt

@@ -179,14 +179,29 @@ internal constructor(
 
   fun onCloseClicked() {
     if (uiState.value is DataCollectionUiState.TaskSubmitted) {
-      viewModelScope.launch { _uiEffects.send(DataCollectionUiEffect.Exit) }
+      exitDataCollection()
     } else {
       showExitWarning()
     }
   }
 
   fun confirmExit() {
     dismissExitWarning()
+    exitDataCollection()
+  }
+
+  fun onBackClicked() {
+    when (val state = _uiState.value) {
+      is DataCollectionUiState.Ready ->
+        if (taskSequenceHandler.isFirstPosition(state.currentTaskId)) showExitWarning()
+        else moveToPreviousTask()
+      is DataCollectionUiState.TaskSubmitted -> exitDataCollection()
+      is DataCollectionUiState.Error,
+      DataCollectionUiState.Loading -> showExitWarning()
+    }
+  }
+
+  private fun exitDataCollection() {
     viewModelScope.launch { _uiEffects.send(DataCollectionUiEffect.Exit) }
   }
 
@@ -229,8 +244,6 @@ internal constructor(
       }
     } ?: false
 
-  fun isAtFirstTask(): Boolean = withReady { taskSequenceHandler.isFirstPosition(it.currentTaskId) }
-
   fun clearDraftBlocking() {
     suppressDrafts()
     clearDraft()

app/src/test/java/org/groundplatform/android/ui/datacollection/DataCollectionFragmentTest.kt

@@ -19,13 +19,16 @@ package org.groundplatform.android.ui.datacollection
 import android.content.Context
 import androidx.compose.ui.test.assertIsDisplayed
 import androidx.compose.ui.test.junit4.createComposeRule
+import androidx.compose.ui.test.onNodeWithContentDescription
 import androidx.compose.ui.test.onNodeWithText
 import androidx.compose.ui.test.performClick
 import androidx.navigation.fragment.findNavController
 import androidx.test.core.app.ApplicationProvider
 import com.google.common.truth.Truth.assertThat
 import dagger.hilt.android.testing.BindValue
 import dagger.hilt.android.testing.HiltAndroidTest
+import javax.inject.Inject
+import kotlin.time.Clock
 import kotlinx.collections.immutable.persistentListOf
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.test.advanceUntilIdle
@@ -70,8 +73,6 @@ import org.junit.runner.RunWith
 import org.mockito.Mock
 import org.robolectric.RobolectricTestRunner
 import org.robolectric.shadows.ShadowToast
-import javax.inject.Inject
-import kotlin.time.Clock
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @HiltAndroidTest
@@ -436,6 +437,48 @@ class DataCollectionFragmentTest : BaseHiltTest() {
       assertNoDraftSaved()
     }
 
+  @Test
+  fun `Back navigation on a later task returns to the previous task`() = runWithTestDispatcher {
+    setupFragment()
+
+    runner()
+      .inputText(TASK_1_RESPONSE)
+      .clickNextButton()
+      .validateTextIsDisplayed(TASK_2_NAME)
+      .pressBackButton()
+      .validateTextIsDisplayed(TASK_1_NAME)
+      .validateTextIsNotDisplayed(TASK_2_NAME)
+  }
+
+  @Test
+  fun `Clicking close button displays the exit confirmation dialog`() = runWithTestDispatcher {
+    setupFragment()
+
+    composeTestRule.onNodeWithContentDescription("Close").performClick()
+    composeTestRule.waitForIdle()
+
+    composeTestRule
+      .onNodeWithText(getString(R.string.data_collection_cancellation_title))
+      .assertIsDisplayed()
+  }
+
+  @Test
+  fun `Clicking close button after submission exits the screen`() = runWithTestDispatcher {
+    setupFragment()
+
+    runner()
+      .inputText(TASK_1_RESPONSE)
+      .clickNextButton()
+      .selectOption(TASK_2_OPTION_LABEL)
+      .clickDoneButton()
+
+    composeTestRule.onNodeWithContentDescription("Close").performClick()
+    advanceUntilIdle()
+
+    assertThat(fragment.findNavController().currentDestination?.id)
+      .isNotEqualTo(R.id.data_collection_fragment)
+  }
+
   @Test
   fun `Multiple choice task remembers previous selection when navigating back and forth`() {
     setupFragment()

app/src/test/java/org/groundplatform/android/ui/datacollection/DataCollectionViewModelTest.kt

@@ -0,0 +1,147 @@
+/*
+ * Copyright 2026 Google LLC
+ *
+ * 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
+ *
+ *     https://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 org.groundplatform.android.ui.datacollection
+
+import androidx.lifecycle.SavedStateHandle
+import app.cash.turbine.test
+import com.google.common.truth.Truth.assertThat
+import dagger.hilt.android.testing.HiltAndroidTest
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.test.advanceUntilIdle
+import org.groundplatform.android.BaseHiltTest
+import org.groundplatform.android.FakeData
+import org.groundplatform.domain.model.task.Task
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.kotlin.any
+import org.mockito.kotlin.anyOrNull
+import org.mockito.kotlin.doReturn
+import org.mockito.kotlin.mock
+import org.mockito.kotlin.whenever
+import org.robolectric.RobolectricTestRunner
+
+@HiltAndroidTest
+@OptIn(ExperimentalCoroutinesApi::class)
+@RunWith(RobolectricTestRunner::class)
+class DataCollectionViewModelTest : BaseHiltTest() {
+
+  @Test
+  fun `onBackClicked shows exit warning while loading`() = runWithTestDispatcher {
+    val viewModel = setupViewModel(DataCollectionUiState.Loading)
+    assertThat(viewModel.uiState.value).isEqualTo(DataCollectionUiState.Loading)
+
+    viewModel.onBackClicked()
+
+    assertThat(viewModel.showExitWarning.value).isTrue()
+  }
+
+  @Test
+  fun `onBackClicked shows exit warning in error state`() = runWithTestDispatcher {
+    val errorState = DataCollectionUiState.Error(DataCollectionErrorCode.SURVEY_LOAD_FAILED)
+    val viewModel = setupViewModel(errorState)
+    advanceUntilIdle()
+    assertThat(viewModel.uiState.value).isEqualTo(errorState)
+
+    viewModel.onBackClicked()
+
+    assertThat(viewModel.showExitWarning.value).isTrue()
+  }
+
+  @Test
+  fun `onBackClicked shows exit warning on the first task`() = runWithTestDispatcher {
+    val viewModel = setupViewModel(READY_STATE)
+    advanceUntilIdle()
+
+    viewModel.onBackClicked()
+
+    assertThat(viewModel.showExitWarning.value).isTrue()
+  }
+
+  @Test
+  fun `onBackClicked emits exit effect when task is submitted`() = runWithTestDispatcher {
+    val viewModel = setupViewModel(DataCollectionUiState.TaskSubmitted(null))
+    advanceUntilIdle()
+
+    viewModel.uiEffects.test {
+      viewModel.onBackClicked()
+      assertThat(awaitItem()).isEqualTo(DataCollectionUiEffect.Exit)
+    }
+  }
+
+  @Test
+  fun `onCloseClicked shows exit warning when not submitted`() = runWithTestDispatcher {
+    val viewModel = setupViewModel(READY_STATE)
+    advanceUntilIdle()
+
+    viewModel.onCloseClicked()
+
+    assertThat(viewModel.showExitWarning.value).isTrue()
+  }
+
+  @Test
+  fun `onCloseClicked emits exit effect when task is submitted`() = runWithTestDispatcher {
+    val viewModel = setupViewModel(DataCollectionUiState.TaskSubmitted(null))
+    advanceUntilIdle()
+
+    viewModel.uiEffects.test {
+      viewModel.onCloseClicked()
+      assertThat(awaitItem()).isEqualTo(DataCollectionUiEffect.Exit)
+    }
+    assertThat(viewModel.showExitWarning.value).isFalse()
+  }
+
+  private suspend fun setupViewModel(result: DataCollectionUiState): DataCollectionViewModel {
+    val initializer =
+      mock<DataCollectionInitializer>().apply {
+        whenever(initialize(any(), any(), anyOrNull(), anyOrNull())) doReturn result
+      }
+
+    return DataCollectionViewModel(
+      savedStateHandle = SavedStateHandle(mapOf("jobId" to "job1")),
+      externalScope = CoroutineScope(testDispatcher),
+      ioDispatcher = testDispatcher,
+      submissionRepository = mock(),
+      submitDataUseCase = mock(),
+      offlineUuidGenerator = mock(),
+      viewModelFactory = mock(),
+      dataCollectionInitializer = initializer,
+      getLoiReportUseCase = mock(),
+    )
+  }
+
+  private companion object {
+    val READY_STATE =
+      DataCollectionUiState.Ready(
+        surveyId = "survey1",
+        job = FakeData.JOB,
+        loiName = "loi",
+        tasks =
+          listOf(
+            Task(
+              id = "taskId",
+              index = 0,
+              type = Task.Type.TEXT,
+              label = "Task",
+              isRequired = false,
+            )
+          ),
+        isAddLoiFlow = false,
+        currentTaskId = "taskId",
+        position = TaskPosition(absoluteIndex = 0, relativeIndex = 0, sequenceSize = 1),
+      )
+  }
+}