Fix FXiOS-12300 [Tab Optimization] Only persist tab data after tabs have been removed

firefox-ios/Client/TabManagement/TabManagerImplementation.swift

@@ -251,16 +251,31 @@ class TabManagerImplementation: NSObject,
                                                 restorePosition: tabs.firstIndex(of: tab),
                                                 isSelected: selectedTab?.tabUUID == tab.tabUUID)
         }
+
+        // Backup tabs for tab undo, this is not a feature on iPhone but is on iPad
         backupCloseTabs = tabs
 
+        // Scroll position for most tabs has been stored via session data when we navigate away,
+        // but we need to save the selected tabs session data to persist scroll position
+        saveSessionData(forTab: selectedTab)
+
         for tab in currentModeTabs {
-            self.removeTab(tab.tabUUID)
+            // Remove each tab without persisting changes
+            self.removeTab(tab, flushToDisk: false)
+            if tab == currentModeTabs.last {
+                // Select tab calls preserve tabs so we don't need to call it again
+                if tab.isPrivate,
+                   let mostRecentActiveTab = mostRecentTab(inTabs: normalActiveTabs) {
+                    // We remove all private tabs so select most recent normal tab
+                    selectTab(mostRecentActiveTab)
+                } else {
+                    // For normal tabs create a new tab and select it
+                    selectTab(addTab())
+                }
+            }
         }
-
         // Save the tab state that existed prior to removals (preserves original selected tab)
         backupCloseTab = currentSelectedTab
-
-        commitChanges()
     }
 
     /// Remove a tab, will notify delegate of the tab removal

firefox-ios/firefox-ios-tests/Tests/ClientTests/TabManagement/TabManagerTests.swift

@@ -104,20 +104,39 @@ class TabManagerTests: XCTestCase {
     @MainActor
     func testRemoveAllTabsCallsSaveTabSession() {
         let subject = createSubject()
-        _ = subject.addTab(URLRequest(url: URL(string: "https://mozilla.com")!), afterTab: nil, isPrivate: false)
+        let tab = subject.addTab(URLRequest(url: URL(string: "https://mozilla.com")!), afterTab: nil, isPrivate: false)
+        subject.selectTab(tab)
         subject.removeAllTabs(isPrivateMode: false)
 
-        XCTAssertEqual(mockSessionStore.saveTabSessionCallCount, 1)
+        // Save tab session is actually called 3 times for one remove all call
+        // 1. Save tab session for currently selected tab before delete to preserve scroll position
+        // 1. AddTab for the new created homepage tab calls commitChanges
+        // 2. selectTab persists changes for currently selected tab before moving to new Tab
+        XCTAssertEqual(mockSessionStore.saveTabSessionCallCount, 3)
     }
 
     @MainActor
-    func testRemoveAllTabsForNotPrivateMode() {
+    func testRemoveAllTabsForNotPrivateModeWhenClosePrivateTabsSettingIsFalse() {
+        (mockProfile.prefs as? MockProfilePrefs)?.things[PrefsKeys.Settings.closePrivateTabs] = false
         var tabs = generateTabs(count: 5)
         tabs.append(contentsOf: generateTabs(ofType: .privateAny, count: 4))
         let subject = createSubject(tabs: tabs)
         XCTAssertEqual(subject.tabs.count, 9)
         subject.removeAllTabs(isPrivateMode: false)
-        XCTAssertEqual(subject.tabs.count, 4)
+        // 5, private mode tabs (4) plus one new normal tab (1)
+        XCTAssertEqual(subject.tabs.count, 5)
+    }
+
+    @MainActor
+    func testRemoveAllTabsForNotPrivateModeWhenClosePrivateTabsSettingIsTrue() {
+        (mockProfile.prefs as? MockProfilePrefs)?.things[PrefsKeys.Settings.closePrivateTabs] = true
+        var tabs = generateTabs(count: 5)
+        tabs.append(contentsOf: generateTabs(ofType: .privateAny, count: 4))
+        let subject = createSubject(tabs: tabs)
+        XCTAssertEqual(subject.tabs.count, 9)
+        subject.removeAllTabs(isPrivateMode: false)
+        // One new normal tab (1)
+        XCTAssertEqual(subject.tabs.count, 1)
     }
 
     @MainActor