[Draw area task] Fix crashes on config change (#2761)

* Get view model from parent fragment

* Revert variable rename for cleaner diff

* Replace uiState with navigationRequests

* Reinit view model on config change

* Remove debugging log

* Move arg key into const

* Log error

* Fix failing check

* Revert changes to manifest

* Revert changes to manifest

* Fix tests

* Add TODOs

* Revert formatting change

* Remove final CR

* Add annotation to dest nav id

* Rename constant

* Inject data collection VM instead of getting from parent

* Remove TODOs

* Fix tests

* Fix failing checks

ground/src/main/AndroidManifest.xml

@@ -98,4 +98,4 @@
 
   </application>
 
-</manifest>
+</manifest>

ground/src/main/java/com/google/android/ground/MainActivity.kt

@@ -90,7 +90,7 @@ class MainActivity : AbstractActivity() {
     viewModel = viewModelFactory[this, MainViewModel::class.java]
 
     lifecycleScope.launch {
-      viewModel.uiState.filterNotNull().collect { updateUi(binding.root, it) }
+      viewModel.navigationRequests.filterNotNull().collect { updateUi(binding.root, it) }
     }
   }
 

ground/src/main/java/com/google/android/ground/MainViewModel.kt

@@ -32,9 +32,9 @@ import com.google.android.ground.ui.common.SharedViewModel
 import com.google.android.ground.util.isPermissionDeniedException
 import javax.inject.Inject
 import kotlinx.coroutines.CoroutineDispatcher
-import kotlinx.coroutines.flow.MutableStateFlow
-import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.SharedFlow
+import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.withContext
 import timber.log.Timber
@@ -53,16 +53,18 @@ constructor(
   authenticationManager: AuthenticationManager,
 ) : AbstractViewModel() {
 
-  private val _uiState: MutableStateFlow<MainUiState?> = MutableStateFlow(null)
-  var uiState: StateFlow<MainUiState?> = _uiState.asStateFlow()
+  private val _navigationRequests: MutableSharedFlow<MainUiState?> = MutableSharedFlow()
+  var navigationRequests: SharedFlow<MainUiState?> = _navigationRequests.asSharedFlow()
 
   /** The window insets determined by the activity. */
   val windowInsets: MutableLiveData<WindowInsetsCompat> = MutableLiveData()
 
   init {
     viewModelScope.launch {
       // TODO: Check auth status whenever fragments resumes
-      authenticationManager.signInState.collect { _uiState.emit(onSignInStateChange(it)) }
+      authenticationManager.signInState.collect {
+        _navigationRequests.emit(onSignInStateChange(it))
+      }
     }
   }
 

ground/src/main/java/com/google/android/ground/ui/common/AbstractMapFragmentWithControls.kt

@@ -47,6 +47,8 @@ abstract class AbstractMapFragmentWithControls : AbstractMapContainerFragment()
     container: ViewGroup?,
     savedInstanceState: Bundle?,
   ): View {
+    super.onCreateView(inflater, container, savedInstanceState)
+
     binding = MapTaskFragBinding.inflate(inflater, container, false)
     binding.fragment = this
     binding.viewModel = getMapViewModel()
@@ -67,13 +69,16 @@ abstract class AbstractMapFragmentWithControls : AbstractMapContainerFragment()
       }
     }
 
+    return binding.root
+  }
+
+  @MustBeInvokedByOverriders
+  override fun onMapReady(map: MapFragment) {
     viewLifecycleOwner.lifecycleScope.launch {
       repeatOnLifecycle(Lifecycle.State.STARTED) {
         getMapViewModel().getCurrentCameraPosition().collect { onMapCameraMoved(it) }
       }
     }
-
-    return binding.root
   }
 
   private fun updateLocationInfoCard(

ground/src/main/java/com/google/android/ground/ui/datacollection/DataCollectionFragment.kt

@@ -51,7 +51,7 @@ class DataCollectionFragment : AbstractFragment(), BackPressListener {
   @Inject lateinit var navigator: Navigator
   @Inject lateinit var viewPagerAdapterFactory: DataCollectionViewPagerAdapterFactory
 
-  private val viewModel: DataCollectionViewModel by hiltNavGraphViewModels(R.id.data_collection)
+  val viewModel: DataCollectionViewModel by hiltNavGraphViewModels(R.id.data_collection)
 
   private lateinit var binding: DataCollectionFragBinding
   private lateinit var progressBar: ProgressBar
@@ -107,7 +107,10 @@ class DataCollectionFragment : AbstractFragment(), BackPressListener {
       }
     )
 
-    lifecycleScope.launch { viewModel.uiState.filterNotNull().collect { updateUI(it) } }
+    lifecycleScope.launch {
+      viewModel.init()
+      viewModel.uiState.filterNotNull().collect { updateUI(it) }
+    }
   }
 
   override fun onResume() {

ground/src/main/java/com/google/android/ground/ui/datacollection/DataCollectionViewModel.kt

@@ -86,7 +86,7 @@ internal constructor(
 ) : AbstractViewModel() {
 
   private val _uiState: MutableStateFlow<UiState?> = MutableStateFlow(null)
-  var uiState: StateFlow<UiState?>
+  var uiState = _uiState.asStateFlow().stateIn(viewModelScope, SharingStarted.Lazily, null)
 
   private val jobId: String = requireNotNull(savedStateHandle[TASK_JOB_ID_KEY])
   private val loiId: String? = savedStateHandle[TASK_LOI_ID_KEY]
@@ -140,15 +140,8 @@ internal constructor(
 
   lateinit var submissionId: String
 
-  init {
-    uiState =
-      _uiState
-        .asStateFlow()
-        .stateIn(
-          viewModelScope,
-          SharingStarted.Lazily,
-          UiState.TaskListAvailable(tasks, getTaskPosition()),
-        )
+  suspend fun init() {
+    _uiState.emit(UiState.TaskListAvailable(tasks, getTaskPosition()))
   }
 
   fun setLoiName(name: String) {

ground/src/main/java/com/google/android/ground/ui/datacollection/DataCollectionViewPagerAdapter.kt

@@ -29,14 +29,18 @@ import com.google.android.ground.ui.datacollection.tasks.text.TextTaskFragment
 import com.google.android.ground.ui.datacollection.tasks.time.TimeTaskFragment
 import dagger.assisted.Assisted
 import dagger.assisted.AssistedInject
+import javax.inject.Provider
 
 /**
  * A simple pager adapter that presents the [Task]s associated with a given Submission, in sequence.
  */
 class DataCollectionViewPagerAdapter
 @AssistedInject
-constructor(@Assisted fragment: Fragment, @Assisted val tasks: List<Task>) :
-  FragmentStateAdapter(fragment) {
+constructor(
+  private val drawAreaTaskFragmentProvider: Provider<DrawAreaTaskFragment>,
+  @Assisted fragment: Fragment,
+  @Assisted val tasks: List<Task>,
+) : FragmentStateAdapter(fragment) {
   override fun getItemCount(): Int = tasks.size
 
   override fun createFragment(position: Int): Fragment {
@@ -48,7 +52,7 @@ constructor(@Assisted fragment: Fragment, @Assisted val tasks: List<Task>) :
         Task.Type.MULTIPLE_CHOICE -> MultipleChoiceTaskFragment()
         Task.Type.PHOTO -> PhotoTaskFragment()
         Task.Type.DROP_PIN -> DropPinTaskFragment()
-        Task.Type.DRAW_AREA -> DrawAreaTaskFragment()
+        Task.Type.DRAW_AREA -> drawAreaTaskFragmentProvider.get()
         Task.Type.NUMBER -> NumberTaskFragment()
         Task.Type.DATE -> DateTaskFragment()
         Task.Type.TIME -> TimeTaskFragment()

ground/src/main/java/com/google/android/ground/ui/datacollection/tasks/polygon/DrawAreaTaskFragment.kt

@@ -15,6 +15,7 @@
  */
 package com.google.android.ground.ui.datacollection.tasks.polygon
 
+import android.os.Bundle
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
@@ -33,18 +34,16 @@ import com.google.android.ground.ui.datacollection.components.TaskView
 import com.google.android.ground.ui.datacollection.components.TaskViewFactory
 import com.google.android.ground.ui.datacollection.tasks.AbstractTaskFragment
 import com.google.android.ground.ui.map.Feature
-import com.google.android.ground.ui.map.MapFragment
 import com.google.android.ground.ui.theme.AppTheme
 import dagger.hilt.android.AndroidEntryPoint
 import javax.inject.Inject
+import javax.inject.Provider
 import kotlinx.coroutines.flow.collectLatest
 import kotlinx.coroutines.launch
 
 @AndroidEntryPoint
-class DrawAreaTaskFragment : AbstractTaskFragment<DrawAreaTaskViewModel>() {
-
-  @Inject lateinit var map: MapFragment
-
+class DrawAreaTaskFragment @Inject constructor() : AbstractTaskFragment<DrawAreaTaskViewModel>() {
+  @Inject lateinit var drawAreaTaskMapFragmentProvider: Provider<DrawAreaTaskMapFragment>
   // Action buttons
   private lateinit var completeButton: TaskButton
   private lateinit var addPointButton: TaskButton
@@ -59,7 +58,10 @@ class DrawAreaTaskFragment : AbstractTaskFragment<DrawAreaTaskViewModel>() {
     // NOTE(#2493): Multiplying by a random prime to allow for some mathematical "uniqueness".
     // Otherwise, the sequentially generated ID might conflict with an ID produced by Google Maps.
     val rowLayout = LinearLayout(requireContext()).apply { id = View.generateViewId() * 11411 }
-    drawAreaTaskMapFragment = DrawAreaTaskMapFragment.newInstance(viewModel, map)
+    drawAreaTaskMapFragment = drawAreaTaskMapFragmentProvider.get()
+    val args = Bundle()
+    args.putString(DrawAreaTaskMapFragment.TASK_ID_FRAGMENT_ARG_KEY, taskId)
+    drawAreaTaskMapFragment.arguments = args
     parentFragmentManager
       .beginTransaction()
       .add(rowLayout.id, drawAreaTaskMapFragment, DrawAreaTaskMapFragment::class.java.simpleName)

ground/src/main/java/com/google/android/ground/ui/datacollection/tasks/polygon/DrawAreaTaskMapFragment.kt

@@ -16,29 +16,43 @@
 package com.google.android.ground.ui.datacollection.tasks.polygon
 
 import android.os.Bundle
+import androidx.hilt.navigation.fragment.hiltNavGraphViewModels
 import androidx.lifecycle.lifecycleScope
+import com.google.android.ground.R
 import com.google.android.ground.ui.common.AbstractMapFragmentWithControls
 import com.google.android.ground.ui.common.BaseMapViewModel
+import com.google.android.ground.ui.datacollection.DataCollectionViewModel
 import com.google.android.ground.ui.map.CameraPosition
 import com.google.android.ground.ui.map.Feature
 import com.google.android.ground.ui.map.MapFragment
 import dagger.hilt.android.AndroidEntryPoint
+import javax.inject.Inject
 import kotlinx.coroutines.launch
 
 @AndroidEntryPoint
-class DrawAreaTaskMapFragment(private val viewModel: DrawAreaTaskViewModel) :
-  AbstractMapFragmentWithControls() {
+class DrawAreaTaskMapFragment @Inject constructor() : AbstractMapFragmentWithControls() {
 
   private lateinit var mapViewModel: BaseMapViewModel
+  private val dataCollectionViewModel: DataCollectionViewModel by
+    hiltNavGraphViewModels(R.id.data_collection)
+  private val viewModel: DrawAreaTaskViewModel by lazy {
+    // Access to this viewModel is lazy for testing. This is because the NavHostController could
+    // not be initialized before the Fragment under test is created, leading to
+    // hiltNavGraphViewModels() to fail when called on launch.
+    val taskId = arguments?.getString(TASK_ID_FRAGMENT_ARG_KEY) ?: error("null taskId fragment arg")
+    dataCollectionViewModel.getTaskViewModel(taskId) as DrawAreaTaskViewModel
+  }
 
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
+
     mapViewModel = getViewModel(BaseMapViewModel::class.java)
   }
 
   override fun getMapViewModel(): BaseMapViewModel = mapViewModel
 
   override fun onMapReady(map: MapFragment) {
+    super.onMapReady(map)
     viewLifecycleOwner.lifecycleScope.launch {
       viewModel.draftArea.collect { feature: Feature? ->
         map.setFeatures(if (feature == null) setOf() else setOf(feature))
@@ -57,7 +71,6 @@ class DrawAreaTaskMapFragment(private val viewModel: DrawAreaTaskViewModel) :
   }
 
   companion object {
-    fun newInstance(viewModel: DrawAreaTaskViewModel, map: MapFragment) =
-      DrawAreaTaskMapFragment(viewModel).apply { this.map = map }
+    const val TASK_ID_FRAGMENT_ARG_KEY = "taskId"
   }
 }

ground/src/test/java/com/google/android/ground/HiltExt.kt

@@ -18,6 +18,7 @@ package com.google.android.ground
 import android.content.ComponentName
 import android.content.Intent
 import android.os.Bundle
+import androidx.annotation.IdRes
 import androidx.annotation.StyleRes
 import androidx.core.os.bundleOf
 import androidx.core.util.Preconditions
@@ -49,7 +50,7 @@ inline fun <reified T : Fragment> launchFragmentInHiltContainer(
 inline fun <reified T : Fragment> launchFragmentWithNavController(
   fragmentArgs: Bundle? = null,
   @StyleRes themeResId: Int = R.style.FragmentScenarioEmptyFragmentActivityTheme,
-  destId: Int,
+  @IdRes destId: Int,
   crossinline preTransactionAction: Fragment.() -> Unit = {},
   crossinline postTransactionAction: Fragment.() -> Unit = {},
 ): ActivityScenario<HiltTestActivity> =
@@ -65,6 +66,7 @@ inline fun <reified T : Fragment> launchFragmentWithNavController(
           navController.setGraph(R.navigation.nav_graph)
           navController.setCurrentDestination(destId, fragmentArgs ?: bundleOf())
 
+          // But we only set the nav controller here.
           // Bind the controller after the view is created but before onViewCreated is called.
           Navigation.setViewNavController(requireView(), navController)
         }
@@ -85,6 +87,8 @@ fun hiltActivityScenario(
         themeResId,
       )
 
+  // Tests fail here with Fragment DrawAreaTaskMapFragment{6a17fdc5}... does not have a
+  // NavController set
   return ActivityScenario.launch(startActivityIntent)
 }