Introduce TaskDataHandler for managing task data state(#2995)

* Create task data handler

* Improve ktdoc

* Add unit tests

* Use google.truth for assertions

* Add null test for getSelection

* Adds test for state flow

* Fix variable names

* Fix variable names

* Add license header

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

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2025 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 com.google.android.ground.ui.datacollection
+
+import com.google.android.ground.model.submission.TaskData
+import com.google.android.ground.model.task.Task
+import com.google.android.ground.model.task.TaskSelections
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+
+/**
+ * Manages the state of [TaskData] associated with [Task] instances.
+ *
+ * This class provides methods to set, retrieve, and observe the data associated with tasks.
+ */
+class TaskDataHandler {
+
+  private val _dataState = MutableStateFlow(LinkedHashMap<Task, TaskData?>())
+
+  val dataState: StateFlow<Map<Task, TaskData?>>
+    get() = _dataState.asStateFlow()
+
+  /**
+   * Sets the [TaskData] for a given [Task].
+   *
+   * If the new value is the same as the current value, no update is performed.
+   *
+   * @param task The [Task] for which to set the data.
+   * @param newValue The new [TaskData] value for the task.
+   */
+  fun setData(task: Task, newValue: TaskData?) {
+    if (getData(task) == newValue) return
+    // Ensure that the map is recreated to ensure that the state flow is emitted.
+    _dataState.value = LinkedHashMap(_dataState.value).apply { this[task] = newValue }
+  }
+
+  /**
+   * Retrieves the [TaskData] associated with a given [Task].
+   *
+   * @param task The [Task] to retrieve data for.
+   * @return The [TaskData] associated with the task, or `null` if no data is found.
+   */
+  fun getData(task: Task): TaskData? = _dataState.value[task]
+
+  /**
+   * Returns a [TaskSelections] map representing the current state of task data.
+   *
+   * This method allows for an optional override of a specific task's value.
+   *
+   * @param taskValueOverride An optional pair of task ID and [TaskData] to override.
+   * @return A [TaskSelections] map containing the current task data, with any specified override
+   *   applied.
+   */
+  fun getTaskSelections(taskValueOverride: Pair<String, TaskData?>? = null): TaskSelections =
+    buildMap {
+      _dataState.value.forEach { (task, value) ->
+        if (taskValueOverride?.first == task.id) {
+            taskValueOverride.second
+          } else {
+            value
+          }
+          ?.apply { put(task.id, this) }
+      }
+    }
+}

ground/src/test/java/com/google/android/ground/ui/datacollection/TaskDataHandlerTest.kt

@@ -0,0 +1,183 @@
+/*
+ * Copyright 2025 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 com.google.android.ground.ui.datacollection
+
+import com.google.android.ground.model.submission.TaskData
+import com.google.android.ground.model.submission.TextTaskData
+import com.google.android.ground.model.task.Task
+import com.google.common.truth.Truth.assertThat
+import com.sharedtest.FakeData
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.flow.toList
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.robolectric.RobolectricTestRunner
+
+@RunWith(RobolectricTestRunner::class)
+class TaskDataHandlerTest {
+
+  @Test
+  fun `setData updates dataState correctly`() = runTest {
+    val handler = TaskDataHandler()
+    val task = createTask("task1")
+    val taskData = createTaskData("data1")
+
+    handler.setData(task, taskData)
+
+    val dataState = handler.dataState.first()
+    assertThat(dataState).hasSize(1)
+    assertThat(dataState[task]).isEqualTo(taskData)
+  }
+
+  @OptIn(ExperimentalCoroutinesApi::class)
+  @Test
+  fun `dataState emits when value is updated`() = runTest {
+    val handler = TaskDataHandler()
+    val task = createTask("task1")
+    val taskData1 = createTaskData("data1")
+    val taskData2 = createTaskData("data2")
+
+    val emissions = mutableListOf<Map<Task, TaskData?>>()
+    val job = launch(UnconfinedTestDispatcher()) { handler.dataState.toList(emissions) }
+
+    handler.setData(task, taskData1)
+    handler.setData(task, taskData2)
+
+    // Verify that both updates were emitted
+    assertThat(emissions).hasSize(3)
+    assertThat(emissions[0]).isEqualTo(emptyMap<Task, TaskData>())
+    assertThat(emissions[1]).isEqualTo(mapOf(task to taskData1))
+    assertThat(emissions[2]).isEqualTo(mapOf(task to taskData2))
+
+    job.cancel()
+  }
+
+  @OptIn(ExperimentalCoroutinesApi::class)
+  @Test
+  fun `dataState does not emit when same value is set`() = runTest {
+    val handler = TaskDataHandler()
+    val task = createTask("task1")
+    val taskData = createTaskData("data1")
+
+    val emissions = mutableListOf<Map<Task, TaskData?>>()
+    val job = launch(UnconfinedTestDispatcher()) { handler.dataState.toList(emissions) }
+
+    handler.setData(task, taskData)
+    handler.setData(task, taskData) // Same value set again
+
+    assertThat(emissions).hasSize(2)
+    assertThat(emissions[0]).isEqualTo(emptyMap<Task, TaskData>())
+    assertThat(emissions[1]).isEqualTo(mapOf(task to taskData))
+
+    job.cancel()
+  }
+
+  @Test
+  fun `getData returns correct data`() = runTest {
+    val handler = TaskDataHandler()
+    val task = createTask("task1")
+    val taskData = createTaskData("data1")
+
+    handler.setData(task, taskData)
+
+    assertThat(handler.getData(task)).isEqualTo(taskData)
+  }
+
+  @Test
+  fun `getData returns null for unknown task`() = runTest {
+    val handler = TaskDataHandler()
+    val task = createTask("task1")
+
+    assertThat(handler.getData(task)).isNull()
+  }
+
+  @Test
+  fun `getTaskSelections returns correct values`() = runTest {
+    val handler = TaskDataHandler()
+    val task1 = createTask("task1")
+    val task2 = createTask("task2")
+    val task3 = createTask("task3")
+    val taskData1 = createTaskData("data1")
+    val taskData2 = createTaskData("data2")
+
+    handler.setData(task1, taskData1)
+    handler.setData(task2, taskData2)
+    handler.setData(task3, null)
+
+    val selections = handler.getTaskSelections()
+    assertThat(selections).hasSize(2)
+    assertThat(selections["task1"]).isEqualTo(taskData1)
+    assertThat(selections["task2"]).isEqualTo(taskData2)
+  }
+
+  @Test
+  fun `getTaskSelections with override returns correct values`() = runTest {
+    val handler = TaskDataHandler()
+    val task1 = createTask("task1")
+    val task2 = createTask("task2")
+    val taskData1 = createTaskData("data1")
+    val taskData2 = createTaskData("data2")
+    val taskDataOverride = createTaskData("override")
+
+    handler.setData(task1, taskData1)
+    handler.setData(task2, taskData2)
+
+    val selections = handler.getTaskSelections(Pair("task1", taskDataOverride))
+    assertThat(selections).hasSize(2)
+    assertThat(selections["task1"]).isEqualTo(taskDataOverride)
+    assertThat(selections["task2"]).isEqualTo(taskData2)
+  }
+
+  @Test
+  fun `getTaskSelections with null override returns correct selections`() = runTest {
+    val handler = TaskDataHandler()
+    val task1 = createTask("task1")
+    val task2 = createTask("task2")
+    val taskData1 = createTaskData("data1")
+    val taskData2 = createTaskData("data2")
+
+    handler.setData(task1, taskData1)
+    handler.setData(task2, taskData2)
+
+    val selections = handler.getTaskSelections(Pair("task1", null))
+    assertThat(selections).hasSize(1)
+    assertThat(selections["task2"]).isEqualTo(taskData2)
+    assertThat(selections["task1"]).isNull()
+  }
+
+  @Test
+  fun `setData with null value`() = runTest {
+    val handler = TaskDataHandler()
+    val task = createTask("task1")
+    val taskData = createTaskData("data1")
+
+    handler.setData(task, taskData)
+    handler.setData(task, null)
+
+    val dataState = handler.dataState.first()
+    assertThat(dataState).hasSize(1)
+    assertThat(dataState[task]).isNull()
+  }
+
+  private fun createTask(taskId: String): Task =
+    FakeData.newTask(id = taskId, type = Task.Type.TEXT)
+
+  private fun createTaskData(value: String): TextTaskData = TextTaskData(value)
+}