Make sure the selected/available rooms are updated when adding space children fails part way through. (#4997)

* Make sure the available rooms are updated if adding space children fails part way through.

* Additionally handle the case where removing rooms fails part way through.

ElementX/Sources/Mocks/SpaceServiceProxyMock.swift

@@ -11,6 +11,10 @@ import Foundation
 import MatrixRustSDK
 import MatrixRustSDKMocks
 
+enum SpaceServiceProxyMockError: Error {
+    case generic
+}
+
 extension SpaceServiceProxyMock {
     struct Configuration {
         var topLevelSpaces: [SpaceServiceRoomProtocol] = []

ElementX/Sources/Screens/Spaces/SpaceAddRoomsScreen/SpaceAddRoomsScreenViewModel.swift

@@ -17,6 +17,9 @@ class SpaceAddRoomsScreenViewModel: SpaceAddRoomsScreenViewModelType, SpaceAddRo
     private let userIndicatorController: UserIndicatorControllerProtocol
     
     private var suggestedRooms: [SpaceAddRoomsScreenRoom] = []
+    /// A place to track rooms that were sucessfully added to the space in order to filter them out from
+    /// the roomsSection when a failure occurs part way through the array.
+    private var addedRoomIDs: Set<String> = []
     
     private var actionsSubject: PassthroughSubject<SpaceAddRoomsScreenViewModelAction, Never> = .init()
     var actions: AnyPublisher<SpaceAddRoomsScreenViewModelAction, Never> {
@@ -89,15 +92,15 @@ class SpaceAddRoomsScreenViewModel: SpaceAddRoomsScreenViewModelType, SpaceAddRo
         }
         
         let existingRooms = spaceRoomListProxy.spaceRoomsPublisher.value
-        var rooms = [SpaceAddRoomsScreenRoom]()
         
-        for summary in roomSummaryProvider.roomListPublisher.value where !summary.isDirect && !existingRooms.contains(where: { $0.id == summary.id }) {
+        let rooms = roomSummaryProvider.roomListPublisher.value
-            rooms.append(.init(summary: summary))
+            .lazy
-        }
+            .filter { [addedRoomIDs] summary in
+                !summary.isDirect && !existingRooms.contains { $0.id == summary.id } && !addedRoomIDs.contains(summary.id)
+            }
+            .map(SpaceAddRoomsScreenRoom.init)
         
-        state.roomsSection = .init(type: .searchResults, rooms: rooms)
+        state.roomsSection = .init(type: .searchResults, rooms: Array(rooms))
-        
-        // rooms crap
     }
     
     /// The actual range values don't matter as long as they contain the lower
@@ -131,15 +134,26 @@ class SpaceAddRoomsScreenViewModel: SpaceAddRoomsScreenViewModelType, SpaceAddRo
         showSavingIndicator()
         defer { hideSavingIndicator() }
         
+        MXLog.info("Adding \(state.selectedRooms.count) rooms to space \(spaceRoomListProxy.id)")
+        
         for room in state.selectedRooms {
-            if case .failure(let error) = await spaceServiceProxy.addChild(room.id, to: spaceRoomListProxy.id) {
+            switch await spaceServiceProxy.addChild(room.id, to: spaceRoomListProxy.id) {
+            case .success:
+                addedRoomIDs.insert(room.id)
+            case .failure(let error):
                 MXLog.error("Failed adding room to space: \(error)")
                 showErrorIndicator()
-                updateRooms() // Hide any rooms that were already added.
+                
+                // Hide rooms that were successfully added.
+                state.selectedRooms.removeAll { addedRoomIDs.contains($0.id) }
+                updateRooms()
+                
                 return
             }
         }
         
+        MXLog.info("\(state.selectedRooms.count) rooms added to space \(spaceRoomListProxy.id)")
+        
         await spaceRoomListProxy.resetAndWaitForFullReload(timeout: .seconds(10))
         
         actionsSubject.send(.dismiss)

ElementX/Sources/Screens/Spaces/SpaceScreen/SpaceScreenModels.swift

@@ -39,6 +39,7 @@ struct SpaceScreenViewState: BindableState {
     
     var editMode: EditMode = .inactive
     var editModeSelectedIDs: Set<String> = []
+    var editModeRemovedIDs: Set<String> = []
     
     var bindings = SpaceScreenViewStateBindings()
     
@@ -50,7 +51,7 @@ struct SpaceScreenViewState: BindableState {
         if editMode == .inactive {
             rooms
         } else {
-            rooms.filter { !$0.isSpace }
+            rooms.filter { !$0.isSpace && !editModeRemovedIDs.contains($0.id) }
         }
     }
     

ElementX/Sources/Screens/Spaces/SpaceScreen/SpaceScreenViewModel.swift

@@ -156,6 +156,7 @@ class SpaceScreenViewModel: SpaceScreenViewModelType, SpaceScreenViewModelProtoc
         case .finishManagingChildren:
             withAnimation(.easeOut(duration: 0.25).disabledDuringTests()) {
                 state.editMode = .inactive
+                state.editModeRemovedIDs = []
             } completion: {
                 self.state.editModeSelectedIDs.removeAll()
             }
@@ -209,16 +210,27 @@ class SpaceScreenViewModel: SpaceScreenViewModelType, SpaceScreenViewModelProtoc
         
         state.bindings.isPresentingRemoveChildrenConfirmation = false
         
+        MXLog.info("Removing \(state.editModeSelectedIDs.count) children from space \(spaceRoomListProxy.id)")
+        
+        var removedIDs: [String] = [] // Using an intermediate array so the screen doesn't change until the operation finishes.
         for childID in state.editModeSelectedIDs {
             switch await spaceServiceProxy.removeChild(childID, from: spaceRoomListProxy.id) {
             case .success:
-                MXLog.info("Successfully removed \(childID) from \(spaceRoomListProxy.id)")
+                removedIDs.append(childID)
-            case .failure:
+            case .failure(let error):
+                MXLog.error("Failed removing room from space: \(error)")
                 showFailureIndicator()
+                
+                // Hide rooms that were successfully removed.
+                state.editModeSelectedIDs = state.editModeSelectedIDs.filter { !removedIDs.contains($0) }
+                state.editModeRemovedIDs.formUnion(removedIDs)
+                
                 return
             }
         }
         
+        MXLog.info("\(state.editModeSelectedIDs.count) children removed from space \(spaceRoomListProxy.id)")
+        
         await spaceRoomListProxy.resetAndWaitForFullReload(timeout: .seconds(10))
         
         process(viewAction: .finishManagingChildren)

UnitTests/Sources/SpaceAddRoomsScreenViewModelTests.swift

@@ -14,6 +14,7 @@ import XCTest
 @MainActor
 class SpaceAddRoomsScreenViewModelTests: XCTestCase {
     var spaceRoomListProxy: SpaceRoomListProxyMock!
+    var spaceServiceProxy: SpaceServiceProxyMock!
     
     var viewModel: SpaceAddRoomsScreenViewModelProtocol!
     var context: SpaceAddRoomsScreenViewModelType.Context { viewModel.context }
@@ -44,15 +45,59 @@ class SpaceAddRoomsScreenViewModelTests: XCTestCase {
         
         try await deferredAction.fulfill()
         
+        XCTAssertTrue(spaceServiceProxy.addChildToCalled, "The room should have been added to the space.")
         XCTAssertTrue(spaceRoomListProxy.resetCalled, "The room list should be reset to pick up the changes.")
     }
     
+    func testFailureWithMultipleRoomsSelected() async throws {
+        // Given a view model with 4 selected rooms.
+        setupViewModel()
+        
+        var deferred = deferFulfillment(context.observe(\.viewState.roomsSection),
+                                        message: "There should be 4 search results.") { section in
+            section.type == .searchResults && section.rooms.count == 4
+        }
+        context.searchQuery = "f"
+        context.send(viewAction: .searchQueryChanged)
+        try await deferred.fulfill()
+        
+        for room in context.viewState.roomsSection.rooms {
+            context.send(viewAction: .toggleRoom(room))
+        }
+        XCTAssertEqual(context.viewState.selectedRooms.count, 4, "All of the rooms should be selected.")
+        
+        // When there's a failure half way through saving.
+        let successfulIDs = context.viewState.roomsSection.rooms.map(\.id).prefix(2)
+        spaceServiceProxy.addChildToClosure = { childID, _ in
+            if successfulIDs.contains(childID) {
+                .success(())
+            } else {
+                .failure(.sdkError(SpaceServiceProxyMockError.generic))
+            }
+        }
+        
+        deferred = deferFulfillment(context.observe(\.viewState.roomsSection),
+                                    message: "The search results should update.") { section in
+            section.type == .searchResults && section.rooms.count == 2
+        }
+        context.send(viewAction: .save)
+        try await deferred.fulfill()
+        
+        // Then the screen should be updated to only show the rooms that still need to be added.
+        XCTAssertEqual(spaceServiceProxy.addChildToCallsCount, 3, "The remaining calls to the service should stop after a failure.")
+        XCTAssertFalse(context.viewState.selectedRooms.contains(where: { successfulIDs.contains($0.id) }),
+                       "The added rooms should no longer show as selected.")
+        XCTAssertFalse(context.viewState.roomsSection.rooms.contains(where: { successfulIDs.contains($0.id) }),
+                       "The added rooms should no longer be listed for selection.")
+    }
+    
     func setupViewModel() {
         let summaryProvider = RoomSummaryProviderMock(.init(state: .loaded(.mockRooms)))
         spaceRoomListProxy = SpaceRoomListProxyMock(.init(spaceServiceRoom: SpaceServiceRoomMock(.init(isSpace: true))))
         
         let clientProxy = ClientProxyMock(.init())
         clientProxy.recentlyVisitedRoomsFilterReturnValue = .init(repeating: JoinedRoomProxyMock(.init()), count: 5)
+        spaceServiceProxy = clientProxy.underlyingSpaceService as? SpaceServiceProxyMock
         
         viewModel = SpaceAddRoomsScreenViewModel(spaceRoomListProxy: spaceRoomListProxy,
                                                  userSession: UserSessionMock(.init(clientProxy: clientProxy)),

UnitTests/Sources/SpaceScreenViewModelTests.swift

@@ -237,6 +237,43 @@ class SpaceScreenViewModelTests: XCTestCase {
         XCTAssertTrue(spaceRoomListProxy.resetCalled, "The room list should be reset to pick up the changes.")
     }
     
+    func testManageRoomsRemovingChildrenWithFailure() async throws {
+        setupViewModel(initialSpaceRooms: mockSpaceRooms)
+        
+        context.send(viewAction: .manageChildren)
+        for room in context.viewState.visibleRooms {
+            context.send(viewAction: .spaceAction(.select(room)))
+        }
+        context.send(viewAction: .removeSelectedChildren)
+        
+        XCTAssertEqual(context.viewState.editMode, .transient, "Managing rooms should enable edit mode.")
+        XCTAssertEqual(context.viewState.visibleRooms.count, 3, "There should be 3 rooms to begin with.")
+        XCTAssertEqual(context.viewState.editModeSelectedIDs.count, 3, "All of the visible rooms should be selected.")
+        XCTAssertTrue(context.isPresentingRemoveChildrenConfirmation, "A confirmation prompt should be shown before removing children.")
+        
+        let successfulIDs = context.viewState.editModeSelectedIDs.prefix(1)
+        spaceServiceProxy.removeChildFromClosure = { childID, _ in
+            if successfulIDs.contains(childID) {
+                .success(())
+            } else {
+                .failure(.sdkError(SpaceServiceProxyMockError.generic))
+            }
+        }
+        
+        let deferred = deferFulfillment(context.observe(\.viewState.visibleRooms.count)) { $0 == 2 }
+        let deferredFailure = deferFailure(context.observe(\.viewState.editMode), timeout: 1) { $0 == .inactive }
+        context.send(viewAction: .confirmRemoveSelectedChildren)
+        try await deferred.fulfill()
+        try await deferredFailure.fulfill()
+        
+        XCTAssertEqual(context.viewState.editMode, .transient, "The screen should remain in edit mode.")
+        XCTAssertEqual(context.viewState.visibleRooms.count, 2, "The removed rooms should no longer be listed for selection.")
+        XCTAssertEqual(context.viewState.editModeSelectedIDs.count, 2, "The removed rooms should no longer be selected.")
+        
+        XCTAssertEqual(spaceServiceProxy.removeChildFromCallsCount, 2, "Each selected room should have been removed.")
+        XCTAssertFalse(spaceRoomListProxy.resetCalled, "The room list should be reset to pick up the changes.")
+    }
+    
     func testLeavingSpace() async throws {
         setupViewModel()
         XCTAssertNil(context.leaveSpaceViewModel)