ViewModel is now access based on taskId using Bundle

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

@@ -46,7 +46,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

ground/src/main/java/com/google/android/ground/ui/datacollection/tasks/location/CaptureLocationTaskFragment.kt

@@ -41,13 +41,10 @@ class CaptureLocationTaskFragment : AbstractTaskFragment<CaptureLocationTaskView
     // 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() * 11149 }
+    val fragment = CaptureLocationTaskMapFragment.newInstance(map, taskId)
     parentFragmentManager
       .beginTransaction()
-      .add(
-        rowLayout.id,
-        CaptureLocationTaskMapFragment.newInstance(viewModel, map),
-        CaptureLocationTaskMapFragment::class.java.simpleName,
-      )
+      .add(rowLayout.id, fragment, CaptureLocationTaskMapFragment::class.java.simpleName)
       .commit()
     return rowLayout
   }

ground/src/main/java/com/google/android/ground/ui/datacollection/tasks/location/CaptureLocationTaskMapFragment.kt

@@ -22,19 +22,33 @@
 import androidx.lifecycle.lifecycleScope
 import com.google.android.ground.ui.common.AbstractMapFragmentWithControls
 import com.google.android.ground.ui.common.BaseMapViewModel
+import com.google.android.ground.ui.datacollection.DataCollectionFragment
 import com.google.android.ground.ui.map.MapFragment
 import dagger.hilt.android.AndroidEntryPoint
 import kotlinx.coroutines.launch
+import timber.log.Timber
 
 @AndroidEntryPoint
-class CaptureLocationTaskMapFragment(private val viewModel: CaptureLocationTaskViewModel) :
-  AbstractMapFragmentWithControls() {
+class CaptureLocationTaskMapFragment : AbstractMapFragmentWithControls() {
 
   private lateinit var mapViewModel: CaptureLocationTaskMapViewModel
+  private lateinit var viewModel: CaptureLocationTaskViewModel
 
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
     mapViewModel = getViewModel(CaptureLocationTaskMapViewModel::class.java)
+    arguments?.let {
+      try {
+        val taskId = it.getString("taskId")
+        taskId?.let {
+          viewModel =
+            (requireParentFragment() as DataCollectionFragment).viewModel.getTaskViewModel(taskId)
+              as CaptureLocationTaskViewModel
+        }
+      } catch (e: Exception) {
+        Timber.e(e)
+      }
+    }
   }
 
   override fun getMapViewModel(): BaseMapViewModel = mapViewModel
@@ -46,18 +60,28 @@
   ): View {
     val root = super.onCreateView(inflater, container, savedInstanceState)
     viewLifecycleOwner.lifecycleScope.launch {
-      getMapViewModel().getLocationUpdates().collect { viewModel.updateLocation(it) }
+      if (this@CaptureLocationTaskMapFragment::viewModel.isInitialized) {
+        getMapViewModel().getLocationUpdates().collect { viewModel.updateLocation(it) }
+      }
     }
     return root
   }
 
   override fun onMapReady(map: MapFragment) {
     binding.locationLockBtn.isClickable = false
-    viewLifecycleOwner.lifecycleScope.launch { viewModel.onMapReady(mapViewModel) }
+    viewLifecycleOwner.lifecycleScope.launch {
+      if (this@CaptureLocationTaskMapFragment::viewModel.isInitialized) {
+        viewModel.onMapReady(mapViewModel)
+      }
+    }
   }
 
   companion object {
-    fun newInstance(viewModel: CaptureLocationTaskViewModel, map: MapFragment) =
-      CaptureLocationTaskMapFragment(viewModel).apply { this.map = map }
+    fun newInstance(map: MapFragment, taskId: String): CaptureLocationTaskMapFragment {
+      val fragment = CaptureLocationTaskMapFragment().apply { this.map = map }
+      val args = Bundle().apply { putString("taskId", taskId) }
+      fragment.arguments = args
+      return fragment
+    }
   }
 }

ground/src/main/java/com/google/android/ground/ui/datacollection/tasks/point/DropPinTaskFragment.kt

@@ -47,9 +47,10 @@ class DropPinTaskFragment : AbstractTaskFragment<DropPinTaskViewModel>() {
     // 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() * 11617 }
+    val fragment = DropPinTaskMapFragment.newInstance(map, taskId)
     parentFragmentManager
       .beginTransaction()
-      .add(rowLayout.id, DropPinTaskMapFragment.newInstance(viewModel, map), "Drop a pin fragment")
+      .add(rowLayout.id, fragment, "Drop a pin fragment")
       .commit()
     return rowLayout
   }

ground/src/main/java/com/google/android/ground/ui/datacollection/tasks/point/DropPinTaskMapFragment.kt

@@ -18,34 +18,56 @@
 import android.os.Bundle
 import com.google.android.ground.ui.common.AbstractMapFragmentWithControls
 import com.google.android.ground.ui.common.BaseMapViewModel
+import com.google.android.ground.ui.datacollection.DataCollectionFragment
 import com.google.android.ground.ui.map.CameraPosition
 import com.google.android.ground.ui.map.MapFragment
 import dagger.hilt.android.AndroidEntryPoint
+import timber.log.Timber
 
 @AndroidEntryPoint
-class DropPinTaskMapFragment(private val viewModel: DropPinTaskViewModel) :
-  AbstractMapFragmentWithControls() {
+class DropPinTaskMapFragment : AbstractMapFragmentWithControls() {
 
   private lateinit var mapViewModel: BaseMapViewModel
+  private lateinit var viewModel: DropPinTaskViewModel
 
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
     mapViewModel = getViewModel(BaseMapViewModel::class.java)
+    arguments?.let {
+      try {
+        val taskId = it.getString("taskId")
+        taskId?.let {
+          viewModel =
+            (requireParentFragment() as DataCollectionFragment).viewModel.getTaskViewModel(taskId)
+              as DropPinTaskViewModel
+        }
+      } catch (e: Exception) {
+        Timber.e(e)
+      }
+    }
   }
 
   override fun getMapViewModel(): BaseMapViewModel = mapViewModel
 
   override fun onMapReady(map: MapFragment) {
-    viewModel.features.observe(this) { map.setFeatures(it) }
+    if (this@DropPinTaskMapFragment::viewModel.isInitialized) {
+      viewModel.features.observe(this) { map.setFeatures(it) }
+    }
   }
 
   override fun onMapCameraMoved(position: CameraPosition) {
     super.onMapCameraMoved(position)
-    viewModel.updateCameraPosition(position)
+    if (this@DropPinTaskMapFragment::viewModel.isInitialized) {
+      viewModel.updateCameraPosition(position)
+    }
   }
 
   companion object {
-    fun newInstance(viewModel: DropPinTaskViewModel, map: MapFragment) =
-      DropPinTaskMapFragment(viewModel).apply { this.map = map }
+    fun newInstance(map: MapFragment, taskId: String): DropPinTaskMapFragment {
+      val fragment = DropPinTaskMapFragment().apply { this.map = map }
+      val args = Bundle().apply { putString("taskId", taskId) }
+      fragment.arguments = args
+      return fragment
+    }
   }
 }

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

@@ -59,7 +59,7 @@ 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 = DrawAreaTaskMapFragment.newInstance(map, taskId)
     parentFragmentManager
       .beginTransaction()
       .add(rowLayout.id, drawAreaTaskMapFragment, DrawAreaTaskMapFragment::class.java.simpleName)

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

@@ -19,36 +19,52 @@
 import androidx.lifecycle.lifecycleScope
 import com.google.android.ground.ui.common.AbstractMapFragmentWithControls
 import com.google.android.ground.ui.common.BaseMapViewModel
+import com.google.android.ground.ui.datacollection.DataCollectionFragment
 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 kotlinx.coroutines.launch
+import timber.log.Timber
 
 @AndroidEntryPoint
-class DrawAreaTaskMapFragment(private val viewModel: DrawAreaTaskViewModel) :
-  AbstractMapFragmentWithControls() {
+class DrawAreaTaskMapFragment : AbstractMapFragmentWithControls() {
 
   private lateinit var mapViewModel: BaseMapViewModel
+  private lateinit var viewModel: DrawAreaTaskViewModel
 
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
     mapViewModel = getViewModel(BaseMapViewModel::class.java)
+    arguments?.let {
+      try {
+        val taskId = it.getString("taskId")
+        taskId?.let {
+          viewModel =
+            (requireParentFragment() as DataCollectionFragment).viewModel.getTaskViewModel(taskId)
+              as DrawAreaTaskViewModel
+        }
+      } catch (e: Exception) {
+        Timber.e(e)
+      }
+    }
   }
 
   override fun getMapViewModel(): BaseMapViewModel = mapViewModel
 
   override fun onMapReady(map: MapFragment) {
     viewLifecycleOwner.lifecycleScope.launch {
-      viewModel.draftArea.collect { feature: Feature? ->
-        map.setFeatures(if (feature == null) setOf() else setOf(feature))
+      if (this@DrawAreaTaskMapFragment::viewModel.isInitialized) {
+        viewModel.draftArea.collect { feature: Feature? ->
+          map.setFeatures(if (feature == null) setOf() else setOf(feature))
+        }
       }
     }
   }
 
   override fun onMapCameraMoved(position: CameraPosition) {
     super.onMapCameraMoved(position)
-    if (!viewModel.isMarkedComplete()) {
+    if (this@DrawAreaTaskMapFragment::viewModel.isInitialized && !viewModel.isMarkedComplete()) {
       val mapCenter = position.coordinates
       viewModel.updateLastVertexAndMaybeCompletePolygon(mapCenter) { c1, c2 ->
         map.getDistanceInPixels(c1, c2)
@@ -57,7 +73,11 @@
   }
 
   companion object {
-    fun newInstance(viewModel: DrawAreaTaskViewModel, map: MapFragment) =
-      DrawAreaTaskMapFragment(viewModel).apply { this.map = map }
+    fun newInstance(map: MapFragment, taskId: String): DrawAreaTaskMapFragment {
+      val fragment = DrawAreaTaskMapFragment().apply { this.map = map }
+      val args = Bundle().apply { putString("taskId", taskId) }
+      fragment.arguments = args
+      return fragment
+    }
   }
 }