Refactor bookmark approach

core/ui/src/main/kotlin/com/google/samples/apps/nowinandroid/core/ui/BookmarkNoteDialog.kt

@@ -18,12 +18,13 @@ import androidx.compose.ui.Modifier
 @Composable
 fun BookmarkNoteDialog(
     initialNote: String = "",
-    onDismiss: (note: String) -> Unit,
+    onDismiss: () -> Unit,
+    onSave: (note: String) -> Unit,
 ) {
     var note by remember { mutableStateOf(initialNote) }
 
     AlertDialog(
-        onDismissRequest = { onDismiss(note) },
+        onDismissRequest = onDismiss,
         title = { Text("Add a note") },
         text = {
             Column {
@@ -42,10 +43,10 @@ fun BookmarkNoteDialog(
             }
         },
         confirmButton = {
-            TextButton(onClick = { onDismiss(note) }) { Text("Save") }
+            TextButton(onClick = { onSave(note) }) { Text("Save") }
         },
         dismissButton = {
-            TextButton(onClick = { onDismiss("") }) { Text("Skip") }
+            TextButton(onClick = onDismiss) { Text("Cancel") }
         },
     )
 }

feature/bookmarks/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksScreen.kt

@@ -106,37 +106,40 @@ internal fun BookmarksScreen(
     modifier: Modifier = Modifier,
     viewModel: BookmarksViewModel = hiltViewModel(),
 ) {
-    val feedState by viewModel.feedUiState.collectAsStateWithLifecycle()
+    val uiState by viewModel.uiState.collectAsStateWithLifecycle()
     var editingNoteId by remember { mutableStateOf<String?>(null) }
 
     BookmarksScreen(
-        feedState = feedState,
+        feedState = uiState.feedState,
         onShowSnackbar = onShowSnackbar,
         removeFromBookmarks = viewModel::removeFromSavedResources,
         onNewsResourceViewed = { viewModel.setNewsResourceViewed(it, true) },
         onTopicClick = onTopicClick,
         modifier = modifier,
-        shouldDisplayUndoBookmark = viewModel.shouldDisplayUndoBookmark,
+        shouldDisplayUndoBookmark = uiState.shouldDisplayUndoBookmark,
         undoBookmarkRemoval = viewModel::undoBookmarkRemoval,
         clearUndoState = viewModel::clearUndoState,
         onEditNote = { editingNoteId = it },
-        isInSelectionMode = viewModel.isInSelectionMode,
-        selectedIds = viewModel.selectedIds,
+        isInSelectionMode = uiState.isInSelectionMode,
+        selectedIds = uiState.selectedIds,
         enterSelectionMode = viewModel::enterSelectionMode,
         exitSelectionMode = viewModel::exitSelectionMode,
         toggleSelection = viewModel::toggleSelection,
         selectAll = viewModel::selectAll,
-        shouldDisplayUndoBulkRemove = viewModel.shouldDisplayUndoBulkRemove,
+        shouldDisplayUndoBulkRemove = uiState.shouldDisplayUndoBulkRemove,
+        bulkRemovedCount = uiState.bulkRemovedCount,
         removeSelected = viewModel::removeSelected,
         undoBulkRemove = viewModel::undoBulkRemove,
         clearBulkUndoState = viewModel::clearBulkUndoState,
     )
 
     editingNoteId?.let { id ->
-        val currentNote = (feedState as? Success)?.feed?.find { it.id == id }?.bookmarkNote ?: ""
+        val currentNote =
+            (uiState.feedState as? Success)?.feed?.find { it.id == id }?.bookmarkNote ?: ""
         BookmarkNoteDialog(
             initialNote = currentNote,
-            onDismiss = { note ->
+            onDismiss = { editingNoteId = null },
+            onSave = { note ->
                 viewModel.updateNote(id, note)
                 editingNoteId = null
             },
@@ -167,13 +170,13 @@ internal fun BookmarksScreen(
     toggleSelection: (String) -> Unit = {},
     selectAll: () -> Unit = {},
     shouldDisplayUndoBulkRemove: Boolean = false,
+    bulkRemovedCount: Int = 0,
     removeSelected: () -> Unit = {},
     undoBulkRemove: () -> Unit = {},
     clearBulkUndoState: () -> Unit = {},
 ) {
     val bookmarkRemovedMessage = stringResource(id = R.string.feature_bookmarks_api_removed)
     val undoText = stringResource(id = R.string.feature_bookmarks_api_undo)
-    val removedCount = remember { androidx.compose.runtime.mutableIntStateOf(0) }
 
     BackHandler(enabled = isInSelectionMode) {
         exitSelectionMode()
@@ -193,7 +196,7 @@ internal fun BookmarksScreen(
     LaunchedEffect(shouldDisplayUndoBulkRemove) {
         if (shouldDisplayUndoBulkRemove) {
             val result = onShowSnackbar(
-                "${removedCount.intValue} bookmarks removed",
+                "$bulkRemovedCount bookmarks removed",
                 undoText,
             )
             if (result) {
@@ -253,10 +256,7 @@ internal fun BookmarksScreen(
                     )
                     TextButton(onClick = selectAll) { Text("All") }
                     Button(
-                        onClick = {
-                            removedCount.intValue = selectedIds.size
-                            removeSelected()
-                        },
+                        onClick = removeSelected,
                         enabled = selectedIds.isNotEmpty(),
                     ) {
                         Text("Remove")

feature/bookmarks/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksUiState.kt

@@ -0,0 +1,29 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * 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.samples.apps.nowinandroid.feature.bookmarks.impl
+
+import com.google.samples.apps.nowinandroid.core.ui.NewsFeedUiState
+import com.google.samples.apps.nowinandroid.core.ui.NewsFeedUiState.Loading
+
+data class BookmarksUiState(
+    val feedState: NewsFeedUiState = Loading,
+    val isInSelectionMode: Boolean = false,
+    val selectedIds: Set<String> = emptySet(),
+    val shouldDisplayUndoBookmark: Boolean = false,
+    val shouldDisplayUndoBulkRemove: Boolean = false,
+    val bulkRemovedCount: Int = 0,
+)

feature/bookmarks/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksViewModel.kt

@@ -16,9 +16,6 @@
 
 package com.google.samples.apps.nowinandroid.feature.bookmarks.impl
 
-import androidx.compose.runtime.getValue
-import androidx.compose.runtime.mutableStateOf
-import androidx.compose.runtime.setValue
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.google.samples.apps.nowinandroid.core.data.repository.UserDataRepository
@@ -27,11 +24,14 @@ import com.google.samples.apps.nowinandroid.core.model.data.UserNewsResource
 import com.google.samples.apps.nowinandroid.core.ui.NewsFeedUiState
 import com.google.samples.apps.nowinandroid.core.ui.NewsFeedUiState.Loading
 import dagger.hilt.android.lifecycle.HiltViewModel
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onStart
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.flow.update
 import kotlinx.coroutines.launch
 import javax.inject.Inject
 
@@ -41,48 +41,48 @@ class BookmarksViewModel @Inject constructor(
     userNewsResourceRepository: UserNewsResourceRepository,
 ) : ViewModel() {
 
-    var shouldDisplayUndoBookmark by mutableStateOf(false)
+    private val _state = MutableStateFlow(BookmarksUiState())
     private var lastRemovedBookmarkId: String? = null
+    private var bulkRemoveSnapshot: List<Pair<String, String?>> = emptyList()
 
-    var isInSelectionMode by mutableStateOf(false)
-        private set
-
-    var selectedIds by mutableStateOf(emptySet<String>())
-        private set
+    val uiState: StateFlow<BookmarksUiState> = combine(
+        userNewsResourceRepository.observeAllBookmarked()
+            .map<List<UserNewsResource>, NewsFeedUiState>(NewsFeedUiState::Success)
+            .onStart { emit(Loading) },
+        _state,
+    ) { feedState, state ->
+        state.copy(feedState = feedState)
+    }.stateIn(
+        scope = viewModelScope,
+        started = SharingStarted.WhileSubscribed(5_000),
+        initialValue = BookmarksUiState(),
+    )
 
     fun enterSelectionMode(initialId: String) {
-        isInSelectionMode = true
-        selectedIds = setOf(initialId)
+        _state.update { it.copy(isInSelectionMode = true, selectedIds = setOf(initialId)) }
     }
 
     fun exitSelectionMode() {
-        isInSelectionMode = false
-        selectedIds = emptySet()
+        _state.update { it.copy(isInSelectionMode = false, selectedIds = emptySet()) }
     }
 
     fun toggleSelection(id: String) {
-        selectedIds = if (id in selectedIds) selectedIds - id else selectedIds + id
+        _state.update { current ->
+            val updated =
+                if (id in current.selectedIds) current.selectedIds - id else current.selectedIds + id
+            current.copy(selectedIds = updated)
+        }
     }
 
     fun selectAll() {
-        val currentFeed = (feedUiState.value as? NewsFeedUiState.Success)?.feed ?: return
-        selectedIds = currentFeed.map { it.id }.toSet()
+        val currentFeed = (uiState.value.feedState as? NewsFeedUiState.Success)?.feed ?: return
+        _state.update { it.copy(selectedIds = currentFeed.map { item -> item.id }.toSet()) }
     }
 
-    val feedUiState: StateFlow<NewsFeedUiState> =
-        userNewsResourceRepository.observeAllBookmarked()
-            .map<List<UserNewsResource>, NewsFeedUiState>(NewsFeedUiState::Success)
-            .onStart { emit(Loading) }
-            .stateIn(
-                scope = viewModelScope,
-                started = SharingStarted.WhileSubscribed(5_000),
-                initialValue = Loading,
-            )
-
     fun removeFromSavedResources(newsResourceId: String) {
+        lastRemovedBookmarkId = newsResourceId
+        _state.update { it.copy(shouldDisplayUndoBookmark = true) }
         viewModelScope.launch {
-            shouldDisplayUndoBookmark = true
-            lastRemovedBookmarkId = newsResourceId
             userDataRepository.setNewsResourceBookmarked(newsResourceId, false)
         }
     }
@@ -103,28 +103,29 @@ class BookmarksViewModel @Inject constructor(
     }
 
     fun clearUndoState() {
-        shouldDisplayUndoBookmark = false
         lastRemovedBookmarkId = null
+        _state.update { it.copy(shouldDisplayUndoBookmark = false) }
     }
 
-    var shouldDisplayUndoBulkRemove by mutableStateOf(false)
-        private set
-
-    private var bulkRemoveSnapshot: List<Pair<String, String?>> = emptyList()
-
     fun removeSelected() {
-        val currentFeed = (feedUiState.value as? NewsFeedUiState.Success)?.feed ?: return
-        bulkRemoveSnapshot = selectedIds.map { id ->
+        val currentFeed = (uiState.value.feedState as? NewsFeedUiState.Success)?.feed ?: return
+        val toRemove = _state.value.selectedIds
+        bulkRemoveSnapshot = toRemove.map { id ->
             id to currentFeed.find { it.id == id }?.bookmarkNote
         }
-        val toRemove = selectedIds.toSet()
         viewModelScope.launch {
             toRemove.forEach { id ->
                 userDataRepository.setNewsResourceBookmarked(id, false)
             }
         }
-        exitSelectionMode()
-        shouldDisplayUndoBulkRemove = true
+        _state.update {
+            it.copy(
+                isInSelectionMode = false,
+                selectedIds = emptySet(),
+                shouldDisplayUndoBulkRemove = true,
+                bulkRemovedCount = toRemove.size,
+            )
+        }
     }
 
     fun undoBulkRemove() {
@@ -140,8 +141,8 @@ class BookmarksViewModel @Inject constructor(
     }
 
     fun clearBulkUndoState() {
-        shouldDisplayUndoBulkRemove = false
         bulkRemoveSnapshot = emptyList()
+        _state.update { it.copy(shouldDisplayUndoBulkRemove = false, bulkRemovedCount = 0) }
     }
 
     fun updateNote(newsResourceId: String, note: String) {

feature/bookmarks/impl/src/test/kotlin/com/google/samples/apps/nowinandroid/feature/bookmarks/impl/BookmarksViewModelTest.kt

@@ -23,7 +23,6 @@ import com.google.samples.apps.nowinandroid.core.testing.repository.TestUserData
 import com.google.samples.apps.nowinandroid.core.testing.util.MainDispatcherRule
 import com.google.samples.apps.nowinandroid.core.ui.NewsFeedUiState.Loading
 import com.google.samples.apps.nowinandroid.core.ui.NewsFeedUiState.Success
-import com.google.samples.apps.nowinandroid.feature.bookmarks.impl.BookmarksViewModel
 import kotlinx.coroutines.flow.collect
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.UnconfinedTestDispatcher
@@ -62,23 +61,23 @@ class BookmarksViewModelTest {
 
     @Test
     fun stateIsInitiallyLoading() = runTest {
-        assertEquals(Loading, viewModel.feedUiState.value)
+        assertEquals(Loading, viewModel.uiState.value.feedState)
     }
 
     @Test
     fun oneBookmark_showsInFeed() = runTest {
-        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.feedUiState.collect() }
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
 
         newsRepository.sendNewsResources(newsResourcesTestData)
         userDataRepository.setNewsResourceBookmarked(newsResourcesTestData[0].id, true)
-        val item = viewModel.feedUiState.value
-        assertIs<Success>(item)
-        assertEquals(item.feed.size, 1)
+        val feedState = viewModel.uiState.value.feedState
+        assertIs<Success>(feedState)
+        assertEquals(feedState.feed.size, 1)
     }
 
     @Test
     fun oneBookmark_whenRemoving_removesFromFeed() = runTest {
-        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.feedUiState.collect() }
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
         // Set the news resources to be used by this test
         newsRepository.sendNewsResources(newsResourcesTestData)
         // Start with the resource saved
@@ -86,59 +85,62 @@ class BookmarksViewModelTest {
         // Use viewModel to remove saved resource
         viewModel.removeFromSavedResources(newsResourcesTestData[0].id)
         // Verify list of saved resources is now empty
-        val item = viewModel.feedUiState.value
-        assertIs<Success>(item)
-        assertEquals(item.feed.size, 0)
-        assertTrue(viewModel.shouldDisplayUndoBookmark)
+        val feedState = viewModel.uiState.value.feedState
+        assertIs<Success>(feedState)
+        assertEquals(feedState.feed.size, 0)
+        assertTrue(viewModel.uiState.value.shouldDisplayUndoBookmark)
     }
 
     @Test
     fun feedUiState_resourceIsViewed_setResourcesViewed() = runTest {
-        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.feedUiState.collect() }
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
 
         // Given
         newsRepository.sendNewsResources(newsResourcesTestData)
         userDataRepository.setNewsResourceBookmarked(newsResourcesTestData[0].id, true)
-        val itemBeforeViewed = viewModel.feedUiState.value
-        assertIs<Success>(itemBeforeViewed)
-        assertFalse(itemBeforeViewed.feed.first().hasBeenViewed)
+        val feedStateBeforeViewed = viewModel.uiState.value.feedState
+        assertIs<Success>(feedStateBeforeViewed)
+        assertFalse(feedStateBeforeViewed.feed.first().hasBeenViewed)
 
         // When
         viewModel.setNewsResourceViewed(newsResourcesTestData[0].id, true)
 
         // Then
-        val item = viewModel.feedUiState.value
-        assertIs<Success>(item)
-        assertTrue(item.feed.first().hasBeenViewed)
+        val feedState = viewModel.uiState.value.feedState
+        assertIs<Success>(feedState)
+        assertTrue(feedState.feed.first().hasBeenViewed)
     }
 
     @Test
     fun enterSelectionMode_setsSelectedId() = runTest {
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
         viewModel.enterSelectionMode("news1")
-        assertTrue(viewModel.isInSelectionMode)
-        assertEquals(setOf("news1"), viewModel.selectedIds)
+        assertTrue(viewModel.uiState.value.isInSelectionMode)
+        assertEquals(setOf("news1"), viewModel.uiState.value.selectedIds)
     }
 
     @Test
     fun toggleSelection_addsAndRemovesId() = runTest {
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
         viewModel.enterSelectionMode("news1")
         viewModel.toggleSelection("news2")
-        assertEquals(setOf("news1", "news2"), viewModel.selectedIds)
+        assertEquals(setOf("news1", "news2"), viewModel.uiState.value.selectedIds)
         viewModel.toggleSelection("news1")
-        assertEquals(setOf("news2"), viewModel.selectedIds)
+        assertEquals(setOf("news2"), viewModel.uiState.value.selectedIds)
     }
 
     @Test
     fun exitSelectionMode_clearsState() = runTest {
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
         viewModel.enterSelectionMode("news1")
         viewModel.exitSelectionMode()
-        assertFalse(viewModel.isInSelectionMode)
-        assertTrue(viewModel.selectedIds.isEmpty())
+        assertFalse(viewModel.uiState.value.isInSelectionMode)
+        assertTrue(viewModel.uiState.value.selectedIds.isEmpty())
     }
 
     @Test
     fun removeSelected_capturesSnapshotBeforeRemoval() = runTest {
-        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.feedUiState.collect() }
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
 
         newsRepository.sendNewsResources(newsResourcesTestData)
         userDataRepository.setNewsResourceBookmarked(newsResourcesTestData[0].id, true)
@@ -149,31 +151,31 @@ class BookmarksViewModelTest {
 
         viewModel.removeSelected()
 
-        assertTrue(viewModel.shouldDisplayUndoBulkRemove)
-        assertFalse(viewModel.isInSelectionMode)
-        assertTrue(viewModel.selectedIds.isEmpty())
+        assertTrue(viewModel.uiState.value.shouldDisplayUndoBulkRemove)
+        assertFalse(viewModel.uiState.value.isInSelectionMode)
+        assertTrue(viewModel.uiState.value.selectedIds.isEmpty())
     }
 
     @Test
     fun feedUiState_undoneBookmarkRemoval_bookmarkIsRestored() = runTest {
-        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.feedUiState.collect() }
+        backgroundScope.launch(UnconfinedTestDispatcher()) { viewModel.uiState.collect() }
 
         // Given
         newsRepository.sendNewsResources(newsResourcesTestData)
         userDataRepository.setNewsResourceBookmarked(newsResourcesTestData[0].id, true)
         viewModel.removeFromSavedResources(newsResourcesTestData[0].id)
-        assertTrue(viewModel.shouldDisplayUndoBookmark)
-        val itemBeforeUndo = viewModel.feedUiState.value
-        assertIs<Success>(itemBeforeUndo)
-        assertEquals(0, itemBeforeUndo.feed.size)
+        assertTrue(viewModel.uiState.value.shouldDisplayUndoBookmark)
+        val feedStateBeforeUndo = viewModel.uiState.value.feedState
+        assertIs<Success>(feedStateBeforeUndo)
+        assertEquals(0, feedStateBeforeUndo.feed.size)
 
         // When
         viewModel.undoBookmarkRemoval()
 
         // Then
-        assertFalse(viewModel.shouldDisplayUndoBookmark)
-        val item = viewModel.feedUiState.value
-        assertIs<Success>(item)
-        assertEquals(1, item.feed.size)
+        assertFalse(viewModel.uiState.value.shouldDisplayUndoBookmark)
+        val feedState = viewModel.uiState.value.feedState
+        assertIs<Success>(feedState)
+        assertEquals(1, feedState.feed.size)
     }
 }

feature/foryou/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/foryou/impl/ForYouScreen.kt

@@ -89,7 +89,6 @@ import com.google.accompanist.permissions.ExperimentalPermissionsApi
 import com.google.accompanist.permissions.PermissionStatus.Denied
 import com.google.accompanist.permissions.rememberPermissionState
 import com.google.samples.apps.nowinandroid.core.designsystem.component.DynamicAsyncImage
-import com.google.samples.apps.nowinandroid.core.ui.BookmarkNoteDialog
 import com.google.samples.apps.nowinandroid.core.designsystem.component.NiaButton
 import com.google.samples.apps.nowinandroid.core.designsystem.component.NiaIconToggleButton
 import com.google.samples.apps.nowinandroid.core.designsystem.component.NiaOverlayLoadingWheel
@@ -100,6 +99,7 @@ import com.google.samples.apps.nowinandroid.core.designsystem.component.scrollba
 import com.google.samples.apps.nowinandroid.core.designsystem.icon.NiaIcons
 import com.google.samples.apps.nowinandroid.core.designsystem.theme.NiaTheme
 import com.google.samples.apps.nowinandroid.core.model.data.UserNewsResource
+import com.google.samples.apps.nowinandroid.core.ui.BookmarkNoteDialog
 import com.google.samples.apps.nowinandroid.core.ui.DevicePreviews
 import com.google.samples.apps.nowinandroid.core.ui.NewsFeedUiState
 import com.google.samples.apps.nowinandroid.core.ui.TrackScreenViewEvent
@@ -120,7 +120,7 @@ fun ForYouScreen(
     val isSyncing by viewModel.isSyncing.collectAsStateWithLifecycle()
     val deepLinkedUserNewsResource by viewModel.deepLinkedNewsResource.collectAsStateWithLifecycle()
 
-    var pendingBookmarkId by remember { mutableStateOf<String?>(null) }
+    var pendingBookmarkNoteId by remember { mutableStateOf<String?>(null) }
 
     ForYouScreen(
         isSyncing = isSyncing,
@@ -132,21 +132,21 @@ fun ForYouScreen(
         onTopicClick = onTopicClick,
         saveFollowedTopics = viewModel::dismissOnboarding,
         onNewsResourcesCheckedChanged = { id, checked ->
+            viewModel.updateNewsResourceSaved(id, checked)
             if (checked) {
-                pendingBookmarkId = id
-            } else {
-                viewModel.updateNewsResourceSaved(id, false)
+                pendingBookmarkNoteId = id
             }
         },
         onNewsResourceViewed = { viewModel.setNewsResourceViewed(it, true) },
         modifier = modifier,
     )
 
-    pendingBookmarkId?.let { id ->
+    pendingBookmarkNoteId?.let { id ->
         BookmarkNoteDialog(
-            onDismiss = { note ->
-                viewModel.bookmarkWithNote(id, note)
-                pendingBookmarkId = null
+            onDismiss = { pendingBookmarkNoteId = null },
+            onSave = { note ->
+                viewModel.updateNote(id, note)
+                pendingBookmarkNoteId = null
             },
         )
     }
@@ -286,7 +286,7 @@ private fun LazyStaggeredGridScope.onboarding(
         OnboardingUiState.Loading,
         OnboardingUiState.LoadFailed,
         OnboardingUiState.NotShown,
-        -> Unit
+            -> Unit
 
         is OnboardingUiState.Shown -> {
             item(span = StaggeredGridItemSpan.FullLine, contentType = "onboarding") {
@@ -515,7 +515,7 @@ private fun feedItemsSize(
         OnboardingUiState.Loading,
         OnboardingUiState.LoadFailed,
         OnboardingUiState.NotShown,
-        -> 0
+            -> 0
 
         is OnboardingUiState.Shown -> 1
     }

feature/foryou/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/foryou/impl/ForYouViewModel.kt

@@ -121,12 +121,9 @@ class ForYouViewModel @Inject constructor(
         }
     }
 
-    fun bookmarkWithNote(newsResourceId: String, note: String) {
+    fun updateNote(newsResourceId: String, note: String) {
         viewModelScope.launch {
-            userDataRepository.setNewsResourceBookmarked(newsResourceId, true)
-            if (note.isNotBlank()) {
-                userDataRepository.setBookmarkNote(newsResourceId, note)
-            }
+            userDataRepository.setBookmarkNote(newsResourceId, note)
         }
     }
 

feature/search/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/search/impl/SearchScreen.kt

@@ -110,7 +110,7 @@ internal fun SearchScreen(
     val searchResultUiState by searchViewModel.searchResultUiState.collectAsStateWithLifecycle()
     val searchQuery by searchViewModel.searchQuery.collectAsStateWithLifecycle()
 
-    var pendingBookmarkId by remember { mutableStateOf<String?>(null) }
+    var pendingBookmarkNoteId by remember { mutableStateOf<String?>(null) }
 
     SearchScreen(
         modifier = modifier,
@@ -122,7 +122,8 @@ internal fun SearchScreen(
         onClearRecentSearches = searchViewModel::clearRecentSearches,
         onNewsResourcesCheckedChanged = { id, checked ->
             if (checked) {
-                pendingBookmarkId = id
+                searchViewModel.setNewsResourceBookmarked(id, true)
+                pendingBookmarkNoteId = id
             } else {
                 searchViewModel.setNewsResourceBookmarked(id, false)
             }
@@ -134,11 +135,12 @@ internal fun SearchScreen(
         onTopicClick = onTopicClick,
     )
 
-    pendingBookmarkId?.let { id ->
+    pendingBookmarkNoteId?.let { id ->
         BookmarkNoteDialog(
-            onDismiss = { note ->
-                searchViewModel.bookmarkWithNote(id, note)
-                pendingBookmarkId = null
+            onDismiss = { pendingBookmarkNoteId = null },
+            onSave = { note ->
+                searchViewModel.updateNote(id, note)
+                pendingBookmarkNoteId = null
             },
         )
     }

feature/search/impl/src/main/kotlin/com/google/samples/apps/nowinandroid/feature/search/impl/SearchViewModel.kt

@@ -121,12 +121,9 @@ class SearchViewModel @Inject constructor(
         }
     }
 
-    fun bookmarkWithNote(newsResourceId: String, note: String) {
+    fun updateNote(newsResourceId: String, note: String) {
         viewModelScope.launch {
-            userDataRepository.setNewsResourceBookmarked(newsResourceId, true)
-            if (note.isNotBlank()) {
-                userDataRepository.setBookmarkNote(newsResourceId, note)
-            }
+            userDataRepository.setBookmarkNote(newsResourceId, note)
         }
     }