7 files changed, 306 insertions, 128 deletions

diff --git a/chrome/browser/sessions/tab_restore_service.cc b/chrome/browser/sessions/tab_restore_service.cc
index d61e42a..6327993 100644
--- a/chrome/browser/sessions/tab_restore_service.cc
+++ b/chrome/browser/sessions/tab_restore_service.cc
@@ -126,17 +126,10 @@ void TabRestoreService::CreateHistoricalTab(NavigationController* tab) {
     return;
 
   scoped_ptr<Tab> local_tab(new Tab());
-  PopulateTabFromController(tab, local_tab.get());
+  PopulateTab(local_tab.get(), browser, tab);
   if (local_tab->navigations.empty())
     return;
 
-  // browser may be NULL when running unit tests.
-  if (browser) {
-    local_tab->browser_id = browser->session_id().id();
-    local_tab->tabstrip_index =
-        browser->tabstrip_model()->GetIndexOfController(tab);
-  }
-
   AddEntry(local_tab.release(), true, true);
 }
 
@@ -152,13 +145,15 @@ void TabRestoreService::BrowserClosing(Browser* browser) {
   window->tabs.resize(browser->tab_count());
   size_t entry_index = 0;
   for (int tab_index = 0; tab_index < browser->tab_count(); ++tab_index) {
-    PopulateTabFromController(
-        &browser->GetTabContentsAt(tab_index)->controller(),
-        &(window->tabs[entry_index]));
-    if (window->tabs[entry_index].navigations.empty())
+    PopulateTab(&(window->tabs[entry_index]),
+                browser,
+                &browser->GetTabContentsAt(tab_index)->controller());
+    if (window->tabs[entry_index].navigations.empty()) {
       window->tabs.erase(window->tabs.begin() + entry_index);
-    else
+    } else {
+      window->tabs[entry_index].browser_id = browser->session_id().id();
       entry_index++;
+    }
   }
   if (window->tabs.empty()) {
     delete window;
@@ -233,12 +228,20 @@ void TabRestoreService::RestoreEntryById(Browser* browser,
         int tab_index = -1;
         if (tab->has_browser())
           tab_browser = BrowserList::FindBrowserWithID(tab->browser_id);
-        if (tab_browser)
+
+        if (tab_browser) {
           tab_index = tab->tabstrip_index;
-        else
-          tab_browser = browser;
-        if (tab_index < 0 || tab_index > browser->tab_count())
-          tab_index = browser->tab_count();
+        } else {
+          tab_browser = Browser::Create(profile());
+          if (tab->has_browser()) {
+            UpdateTabBrowserIDs(tab->browser_id,
+                                tab_browser->session_id().id());
+          }
+          tab_browser->window()->Show();
+        }
+
+        if (tab_index < 0 || tab_index > tab_browser->tab_count())
+          tab_index = tab_browser->tab_count();
         tab_browser->AddRestoredTab(tab->navigations, tab_index,
                                     tab->current_navigation_index, true);
       }
@@ -255,6 +258,11 @@ void TabRestoreService::RestoreEntryById(Browser* browser,
         if (restored_tab)
           restored_tab->controller().LoadIfNecessary();
       }
+      // All the window's tabs had the same former browser_id.
+      if (window->tabs[0].has_browser()) {
+        UpdateTabBrowserIDs(window->tabs[0].browser_id,
+                            browser->session_id().id());
+      }
       browser->window()->Show();
     } else {
       NOTREACHED();
@@ -325,9 +333,9 @@ void TabRestoreService::Save() {
   BaseSessionService::Save();
 }
 
-void TabRestoreService::PopulateTabFromController(
-    NavigationController* controller,
-    Tab* tab) {
+void TabRestoreService::PopulateTab(Tab* tab,
+                                    Browser* browser,
+                                    NavigationController* controller) {
   const int pending_index = controller->pending_entry_index();
   int entry_count = controller->entry_count();
   if (entry_count == 0 && pending_index == 0)
@@ -341,6 +349,13 @@ void TabRestoreService::PopulateTabFromController(
   tab->current_navigation_index = controller->GetCurrentEntryIndex();
   if (tab->current_navigation_index == -1 && entry_count > 0)
     tab->current_navigation_index = 0;
+
+  // Browser may be NULL during unit tests.
+  if (browser) {
+    tab->browser_id = browser->session_id().id();
+    tab->tabstrip_index =
+        browser->tabstrip_model()->GetIndexOfController(controller);
+  }
 }
 
 void TabRestoreService::NotifyTabsChanged() {
@@ -683,6 +698,18 @@ void TabRestoreService::ValidateAndDeleteEmptyEntries(
   STLDeleteElements(&invalid_entries);
 }
 
+void TabRestoreService::UpdateTabBrowserIDs(SessionID::id_type old_id,
+                                            SessionID::id_type new_id) {
+  for (Entries::iterator i = entries_.begin(); i != entries_.end(); ++i) {
+    Entry* entry = *i;
+    if (entry->type == TAB) {
+      Tab* tab = static_cast<Tab*>(entry);
+      if (tab->browser_id == old_id)
+        tab->browser_id = new_id;
+    }
+  }
+}
+
 void TabRestoreService::OnGotPreviousSession(
     Handle handle,
     std::vector<SessionWindow*>* windows) {
diff --git a/chrome/browser/sessions/tab_restore_service.h b/chrome/browser/sessions/tab_restore_service.h
index c728652..6b1da34 100644
--- a/chrome/browser/sessions/tab_restore_service.h
+++ b/chrome/browser/sessions/tab_restore_service.h
@@ -168,9 +168,11 @@ class TabRestoreService : public BaseSessionService {
     LOADED_LAST_SESSION  = 1 << 4
   };
 
-  // Populates tabs->navigations from the NavigationController.
-  void PopulateTabFromController(NavigationController* controller,
-                                 Tab* tab);
+  // Populates the tab's navigations from the NavigationController, and its
+  // browser_id and tabstrip_index from the browser.
+  void PopulateTab(Tab* tab,
+                   Browser* browser,
+                   NavigationController* controller);
 
   // Notifies observers the tabs have changed.
   void NotifyTabsChanged();
@@ -235,6 +237,10 @@ class TabRestoreService : public BaseSessionService {
   // hold.
   void ValidateAndDeleteEmptyEntries(std::vector<Entry*>* entries);
 
+  // Finds tab entries with the old browser_id and sets it to the new one.
+  void UpdateTabBrowserIDs(SessionID::id_type old_id,
+                           SessionID::id_type new_id);
+
   // Callback from SessionService when we've received the windows from the
   // previous session. This creates and add entries to |staging_entries_|
   // and invokes LoadStateChanged.
diff --git a/chrome/browser/tab_restore_uitest.cc b/chrome/browser/tab_restore_uitest.cc
index 7c7956f..af98b5c 100644
--- a/chrome/browser/tab_restore_uitest.cc
+++ b/chrome/browser/tab_restore_uitest.cc
@@ -28,26 +28,76 @@ class TabRestoreUITest : public UITest {
   }
 
  protected:
-  void RestoreTab() {
-    int tab_count;
-
-    // Reset browser_proxy to new window.
-    scoped_ptr<BrowserProxy> browser_proxy(automation()->GetBrowserWindow(0));
+  // Uses the undo-close-tab accelerator to undo a close-tab or close-window
+  // operation. The newly restored tab is expected to appear in the
+  // window at index |expected_window_index|, at the |expected_tabstrip_index|,
+  // and to be active. If |expected_window_index| is equal to the number of
+  // current windows, the restored tab is expected to be created in a new
+  // window (since the index is 0-based).
+  void RestoreTab(int expected_window_index,
+                  int expected_tabstrip_index) {
+    int tab_count = 0;
+    int window_count = 0;
+
+    ASSERT_TRUE(automation()->GetBrowserWindowCount(&window_count));
+    ASSERT_GT(window_count, 0);
+
+    bool expect_new_window = (expected_window_index == window_count);
+    scoped_ptr<BrowserProxy> browser_proxy;
+    if (expect_new_window) {
+      browser_proxy.reset(automation()->GetBrowserWindow(0));
+    } else {
+      ASSERT_GT(window_count, expected_window_index);
+      browser_proxy.reset(
+          automation()->GetBrowserWindow(expected_window_index));
+    }
     ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
     ASSERT_GT(tab_count, 0);
 
     // Restore the tab.
     ASSERT_TRUE(browser_proxy->ApplyAccelerator(IDC_RESTORE_TAB));
-    ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(
-        tab_count + 1, action_max_timeout_ms()));
+
+    if (expect_new_window) {
+      ASSERT_TRUE(automation()->WaitForWindowCountToBecome(
+          ++window_count, action_max_timeout_ms()));
+      browser_proxy.reset(automation()->
+          GetBrowserWindow(expected_window_index));
+    } else {
+      ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(
+          ++tab_count, action_max_timeout_ms()));
+    }
 
     // Get a handle to the restored tab.
     ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
+    ASSERT_GT(tab_count, expected_tabstrip_index);
     scoped_ptr<TabProxy> restored_tab_proxy(
-        browser_proxy->GetTab(tab_count - 1));
+        browser_proxy->GetTab(expected_tabstrip_index));
     // Wait for the restored tab to finish loading.
     ASSERT_TRUE(restored_tab_proxy->WaitForTabToBeRestored(
         action_max_timeout_ms()));
+
+    // Ensure that the tab and window are active.
+    CheckActiveWindow(browser_proxy.get());
+    EXPECT_EQ(expected_tabstrip_index,
+        GetActiveTabIndex(expected_window_index));
+  }
+
+  // Adds tabs to the given browser, all navigated to url1_. Returns
+  // the final number of tabs.
+  int AddSomeTabs(BrowserProxy* browser, int how_many) {
+    int starting_tab_count = -1;
+    // Use EXPECT instead of ASSERT throughout to avoid trying to return void.
+    EXPECT_TRUE(browser->GetTabCount(&starting_tab_count));
+
+    for (int i = 0; i < how_many; ++i) {
+      browser->AppendTab(url1_);
+      EXPECT_TRUE(browser->WaitForTabCountToBecome(starting_tab_count + i + 1,
+                                                   action_max_timeout_ms()));
+    }
+    int tab_count;
+    EXPECT_TRUE(browser->GetTabCount(&tab_count));
+    EXPECT_EQ(starting_tab_count + how_many, tab_count);
+    return tab_count;
   }
 
   // Ensure that the given browser occupies the currently active window.
@@ -87,16 +137,9 @@ class TabRestoreUITest : public UITest {
 TEST_F(TabRestoreUITest, Basic) {
   scoped_ptr<BrowserProxy> browser_proxy(automation()->GetBrowserWindow(0));
 
-  int tab_count;
-  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  int starting_tab_count = tab_count;
-
-  // Add a tab
-  browser_proxy->AppendTab(url1_);
-  ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(tab_count + 1,
-                                                     action_max_timeout_ms()));
-  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count + 1, tab_count);
+  int starting_tab_count;
+  ASSERT_TRUE(browser_proxy->GetTabCount(&starting_tab_count));
+  int tab_count = AddSomeTabs(browser_proxy.get(), 1);
 
   int closed_tab_index = tab_count - 1;
   scoped_ptr<TabProxy> new_tab(browser_proxy->GetTab(closed_tab_index));
@@ -106,15 +149,15 @@ TEST_F(TabRestoreUITest, Basic) {
   new_tab->Close(true);
   new_tab.reset();
   ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count, tab_count);
+  EXPECT_EQ(starting_tab_count, tab_count);
 
-  RestoreTab();
+  RestoreTab(0, closed_tab_index);
 
   // And make sure everything looks right.
   ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count + 1, tab_count);
-  ASSERT_EQ(closed_tab_index, GetActiveTabIndex());
-  ASSERT_EQ(url1_, GetActiveTabURL());
+  EXPECT_EQ(starting_tab_count + 1, tab_count);
+  EXPECT_EQ(closed_tab_index, GetActiveTabIndex());
+  EXPECT_EQ(url1_, GetActiveTabURL());
 }
 
 // Close a tab not at the end of the current window, then restore it. The tab
@@ -122,22 +165,9 @@ TEST_F(TabRestoreUITest, Basic) {
 TEST_F(TabRestoreUITest, MiddleTab) {
   scoped_ptr<BrowserProxy> browser_proxy(automation()->GetBrowserWindow(0));
 
-  int tab_count;
-  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  int starting_tab_count = tab_count;
-
-  // Add a couple of tabs
-  browser_proxy->AppendTab(url1_);
-  ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(tab_count + 1,
-                                                     action_max_timeout_ms()));
-  browser_proxy->AppendTab(url1_);
-  ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(tab_count + 2,
-                                                     action_max_timeout_ms()));
-  browser_proxy->AppendTab(url1_);
-  ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(tab_count + 3,
-                                                     action_max_timeout_ms()));
-  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count + 3, tab_count);
+  int starting_tab_count;
+  ASSERT_TRUE(browser_proxy->GetTabCount(&starting_tab_count));
+  int tab_count = AddSomeTabs(browser_proxy.get(), 3);
 
   // Close one in the middle
   int closed_tab_index = starting_tab_count + 1;
@@ -148,15 +178,15 @@ TEST_F(TabRestoreUITest, MiddleTab) {
   new_tab->Close(true);
   new_tab.reset();
   ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count + 2, tab_count);
+  EXPECT_EQ(starting_tab_count + 2, tab_count);
 
-  RestoreTab();
+  RestoreTab(0, closed_tab_index);
 
   // And make sure everything looks right.
   ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count + 3, tab_count);
-  ASSERT_EQ(closed_tab_index, GetActiveTabIndex());
-  ASSERT_EQ(url1_, GetActiveTabURL());
+  EXPECT_EQ(starting_tab_count + 3, tab_count);
+  EXPECT_EQ(closed_tab_index, GetActiveTabIndex());
+  EXPECT_EQ(url1_, GetActiveTabURL());
 }
 
 // Close a tab, switch windows, then restore the tab. The tab should be in its
@@ -169,22 +199,9 @@ TEST_F(TabRestoreUITest, RestoreToDifferentWindow) {
   // CheckActiveWindow(). See comments in that function.
   CheckActiveWindow(browser_proxy.get());
 
-  int tab_count;
-  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  int starting_tab_count = tab_count;
-
-  // Add a couple of tabs
-  browser_proxy->AppendTab(url1_);
-  ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(tab_count + 1,
-                                                     action_max_timeout_ms()));
-  browser_proxy->AppendTab(url1_);
-  ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(tab_count + 2,
-                                                     action_max_timeout_ms()));
-  browser_proxy->AppendTab(url1_);
-  ASSERT_TRUE(browser_proxy->WaitForTabCountToBecome(tab_count + 3,
-                                                     action_max_timeout_ms()));
-  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count + 3, tab_count);
+  int starting_tab_count;
+  ASSERT_TRUE(browser_proxy->GetTabCount(&starting_tab_count));
+  int tab_count = AddSomeTabs(browser_proxy.get(), 3);
 
   // Close one in the middle
   int closed_tab_index = starting_tab_count + 1;
@@ -195,7 +212,7 @@ TEST_F(TabRestoreUITest, RestoreToDifferentWindow) {
   new_tab->Close(true);
   new_tab.reset();
   ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
-  ASSERT_EQ(starting_tab_count + 2, tab_count);
+  EXPECT_EQ(starting_tab_count + 2, tab_count);
 
   // Create a new browser.
   ASSERT_TRUE(automation()->OpenNewBrowserWindow(false));
@@ -204,19 +221,20 @@ TEST_F(TabRestoreUITest, RestoreToDifferentWindow) {
 
   CheckActiveWindow(automation()->GetBrowserWindow(1));
 
-  RestoreTab();
+  // Restore tab into original browser.
+  RestoreTab(0, closed_tab_index);
 
   // And make sure everything looks right.
   CheckActiveWindow(browser_proxy.get());
   ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
   EXPECT_EQ(starting_tab_count + 3, tab_count);
-  ASSERT_EQ(closed_tab_index, GetActiveTabIndex());
-  ASSERT_EQ(url1_, GetActiveTabURL());
+  EXPECT_EQ(closed_tab_index, GetActiveTabIndex(0));
+  EXPECT_EQ(url1_, GetActiveTabURL(0));
 }
 
 // Close a tab, open a new window, close the first window, then restore the
-// tab. It should be at the end of the current (new) window's tabstrip.
-TEST_F(TabRestoreUITest, RestoreFromClosedWindow) {
+// tab. It should be in a new window.
+TEST_F(TabRestoreUITest, BasicRestoreFromClosedWindow) {
   // This test is disabled on win2k. See bug 1215881.
   if (win_util::GetWinVersion() == win_util::WINVERSION_2000)
     return;
@@ -245,7 +263,7 @@ TEST_F(TabRestoreUITest, RestoreFromClosedWindow) {
       2, action_max_timeout_ms()));
   CheckActiveWindow(automation()->GetBrowserWindow(1));
 
-  // Close the first browser.
+  // Close the final tab in the first browser.
   EXPECT_TRUE(tab_proxy->Close(true));
   ASSERT_TRUE(automation()->WaitForWindowCountToBecome(
       1, action_max_timeout_ms()));
@@ -254,16 +272,120 @@ TEST_F(TabRestoreUITest, RestoreFromClosedWindow) {
   tab_proxy.reset();
   browser_proxy.reset();
 
-  RestoreTab();
+  RestoreTab(1, 0);
 
-  // Tab should be at the end of the current (only) window.
-  browser_proxy.reset(automation()->GetBrowserWindow(0));
+  // Tab should be in a new window.
+  browser_proxy.reset(automation()->GetBrowserWindow(1));
   CheckActiveWindow(browser_proxy.get());
   tab_proxy.reset(browser_proxy->GetActiveTab());
   // And make sure the URLs matches.
-  ASSERT_EQ(url2_, GetActiveTabURL());
-  ASSERT_TRUE(tab_proxy->GoBack());
-  ASSERT_EQ(url1_, GetActiveTabURL());
+  EXPECT_EQ(url2_, GetActiveTabURL(1));
+  EXPECT_TRUE(tab_proxy->GoBack());
+  EXPECT_EQ(url1_, GetActiveTabURL(1));
+}
+
+// Open a window with multiple tabs, close a tab, then close the window.
+// Restore both and make sure the tab goes back into the window.
+TEST_F(TabRestoreUITest, RestoreWindowAndTab) {
+  scoped_ptr<BrowserProxy> browser_proxy(automation()->GetBrowserWindow(0));
+  CheckActiveWindow(browser_proxy.get());
+
+  int starting_tab_count;
+  ASSERT_TRUE(browser_proxy->GetTabCount(&starting_tab_count));
+  int tab_count = AddSomeTabs(browser_proxy.get(), 3);
+
+  // Close one in the middle
+  int closed_tab_index = starting_tab_count + 1;
+  scoped_ptr<TabProxy> new_tab(browser_proxy->GetTab(closed_tab_index));
+  // Make sure we're at url.
+  new_tab->NavigateToURL(url1_);
+  // Close the tab.
+  new_tab->Close(true);
+  new_tab.reset();
+  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
+  EXPECT_EQ(starting_tab_count + 2, tab_count);
+
+  // Create a new browser.
+  ASSERT_TRUE(automation()->OpenNewBrowserWindow(false));
+  ASSERT_TRUE(automation()->WaitForWindowCountToBecome(
+      2, action_max_timeout_ms()));
+  CheckActiveWindow(automation()->GetBrowserWindow(1));
+
+  // Close the first browser.
+  bool application_closing;
+  EXPECT_TRUE(CloseBrowser(browser_proxy.get(), &application_closing));
+  EXPECT_FALSE(application_closing);
+  ASSERT_TRUE(automation()->WaitForWindowCountToBecome(
+      1, action_max_timeout_ms()));
+
+  // Browser is no longer valid.
+  browser_proxy.reset();
+
+  // Restore the first window. The expected_tabstrip_index (second argument)
+  // indicates the expected active tab.
+  RestoreTab(1, starting_tab_count + 1);
+  browser_proxy.reset(automation()->GetBrowserWindow(1));
+  CheckActiveWindow(browser_proxy.get());
+  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
+  EXPECT_EQ(starting_tab_count + 2, tab_count);
+
+  // Restore the closed tab.
+  RestoreTab(1, closed_tab_index);
+  CheckActiveWindow(browser_proxy.get());
+  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
+  EXPECT_EQ(starting_tab_count + 3, tab_count);
+  EXPECT_EQ(url1_, GetActiveTabURL(1));
+}
+
+// Open a window with two tabs, close both (closing the window), then restore
+// both. Make sure both restored tabs are in the same window.
+TEST_F(TabRestoreUITest, RestoreIntoSameWindow) {
+  scoped_ptr<BrowserProxy> browser_proxy(automation()->GetBrowserWindow(0));
+  CheckActiveWindow(browser_proxy.get());
+
+  int starting_tab_count;
+  ASSERT_TRUE(browser_proxy->GetTabCount(&starting_tab_count));
+  int tab_count = AddSomeTabs(browser_proxy.get(), 2);
+
+  // Navigate the rightmost one to url2_ for easier identification.
+  scoped_ptr<TabProxy> tab_proxy(browser_proxy->GetTab(tab_count - 1));
+  tab_proxy->NavigateToURL(url2_);
+
+  // Create a new browser.
+  ASSERT_TRUE(automation()->OpenNewBrowserWindow(false));
+  ASSERT_TRUE(automation()->WaitForWindowCountToBecome(
+      2, action_max_timeout_ms()));
+  CheckActiveWindow(automation()->GetBrowserWindow(1));
+
+  // Close all but one tab in the first browser, left to right.
+  while (tab_count > 1) {
+    scoped_ptr<TabProxy> tab_to_close(browser_proxy->GetTab(0));
+    tab_to_close->Close(true);
+    ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
+  }
+
+  // Close the last tab, closing the browser.
+  tab_proxy.reset(browser_proxy->GetTab(0));
+  EXPECT_TRUE(tab_proxy->Close(true));
+  ASSERT_TRUE(automation()->WaitForWindowCountToBecome(
+      1, action_max_timeout_ms()));
+  browser_proxy.reset();
+  tab_proxy.reset();
+
+  // Restore the last-closed tab into a new window.
+  RestoreTab(1, 0);
+  browser_proxy.reset(automation()->GetBrowserWindow(1));
+  CheckActiveWindow(browser_proxy.get());
+  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
+  EXPECT_EQ(1, tab_count);
+  EXPECT_EQ(url2_, GetActiveTabURL(1));
+
+  // Restore the next-to-last-closed tab into the same window.
+  RestoreTab(1, 0);
+  CheckActiveWindow(browser_proxy.get());
+  ASSERT_TRUE(browser_proxy->GetTabCount(&tab_count));
+  EXPECT_EQ(2, tab_count);
+  EXPECT_EQ(url1_, GetActiveTabURL(1));
 }
 
 // Tests that a duplicate history entry is not created when we restore a page
@@ -300,13 +422,13 @@ TEST_F(TabRestoreUITest, RestoreWithExistingSiteInstance) {
   browser_proxy->AppendTab(http_url2);
 
   // Restore the closed tab.
-  RestoreTab();
+  RestoreTab(0, tab_count - 1);
   tab.reset(browser_proxy->GetActiveTab());
 
   // And make sure the URLs match.
-  ASSERT_EQ(http_url2, GetActiveTabURL());
-  ASSERT_TRUE(tab->GoBack());
-  ASSERT_EQ(http_url1, GetActiveTabURL());
+  EXPECT_EQ(http_url2, GetActiveTabURL());
+  EXPECT_TRUE(tab->GoBack());
+  EXPECT_EQ(http_url1, GetActiveTabURL());
 }
 
 // Tests that the SiteInstances used for entries in a restored tab's history
@@ -346,19 +468,19 @@ TEST_F(TabRestoreUITest, RestoreCrossSiteWithExistingSiteInstance) {
   browser_proxy->AppendTab(http_url2);
 
   // Restore the closed tab.
-  RestoreTab();
+  RestoreTab(0, tab_count - 1);
   tab.reset(browser_proxy->GetActiveTab());
 
   // And make sure the URLs match.
-  ASSERT_EQ(url1_, GetActiveTabURL());
-  ASSERT_TRUE(tab->GoBack());
-  ASSERT_EQ(http_url1, GetActiveTabURL());
+  EXPECT_EQ(url1_, GetActiveTabURL());
+  EXPECT_TRUE(tab->GoBack());
+  EXPECT_EQ(http_url1, GetActiveTabURL());
 
   // Navigating to a new URL should clear the forward list, because the max
   // page ID of the renderer should have been updated when we restored the tab.
   tab->NavigateToURL(http_url2);
-  ASSERT_FALSE(tab->GoForward());
-  ASSERT_EQ(http_url2, GetActiveTabURL());
+  EXPECT_FALSE(tab->GoForward());
+  EXPECT_EQ(http_url2, GetActiveTabURL());
 }
 
 TEST_F(TabRestoreUITest, RestoreWindow) {
@@ -406,11 +528,11 @@ TEST_F(TabRestoreUITest, RestoreWindow) {
   ASSERT_TRUE(restored_tab_proxy->WaitForTabToBeRestored(action_timeout_ms()));
   GURL url;
   ASSERT_TRUE(restored_tab_proxy->GetCurrentURL(&url));
-  ASSERT_TRUE(url == url1_);
+  EXPECT_TRUE(url == url1_);
 
   restored_tab_proxy.reset(
         browser_proxy->GetTab(initial_tab_count + 1));
   ASSERT_TRUE(restored_tab_proxy->WaitForTabToBeRestored(action_timeout_ms()));
   ASSERT_TRUE(restored_tab_proxy->GetCurrentURL(&url));
-  ASSERT_TRUE(url == url2_);
+  EXPECT_TRUE(url == url2_);
 }
diff --git a/chrome/test/automation/browser_proxy.cc b/chrome/test/automation/browser_proxy.cc
index d8240a0..03a1e51 100644
--- a/chrome/test/automation/browser_proxy.cc
+++ b/chrome/test/automation/browser_proxy.cc
@@ -19,6 +19,7 @@
 using base::TimeDelta;
 using base::TimeTicks;
 
+
 bool BrowserProxy::ActivateTab(int tab_index) {
   return ActivateTabWithTimeout(tab_index, base::kNoTimeout, NULL);
 }
diff --git a/chrome/test/automation/browser_proxy.h b/chrome/test/automation/browser_proxy.h
index 68d75bae..7f69681 100644
--- a/chrome/test/automation/browser_proxy.h
+++ b/chrome/test/automation/browser_proxy.h
@@ -110,7 +110,6 @@ class BrowserProxy : public AutomationResourceProxy {
   // On failure, returns NULL.
   AutocompleteEditProxy* GetAutocompleteEdit();
 
-
   // Apply the accelerator with given id (IDC_BACK, IDC_NEWTAB ...)
   // The list can be found at chrome/app/chrome_dll_resource.h
   // Returns true if the call was successful.
diff --git a/chrome/test/ui/ui_test.cc b/chrome/test/ui/ui_test.cc
index 620637e..d5ac11d 100644
--- a/chrome/test/ui/ui_test.cc
+++ b/chrome/test/ui/ui_test.cc
@@ -458,8 +458,14 @@ void UITest::CleanupAppProcesses() {
 #endif
 }
 
-TabProxy* UITest::GetActiveTab() {
-  scoped_ptr<BrowserProxy> window_proxy(automation()->GetBrowserWindow(0));
+TabProxy* UITest::GetActiveTab(int window_index) {
+  EXPECT_GE(window_index, 0);
+  int window_count;
+  // Use EXPECT rather than ASSERT here because ASSERT_* returns void.
+  EXPECT_TRUE(automation()->GetBrowserWindowCount(&window_count));
+  EXPECT_GT(window_count, window_index);
+  scoped_ptr<BrowserProxy> window_proxy(automation()->
+      GetBrowserWindow(window_index));
   if (!window_proxy.get())
     return NULL;
 
@@ -545,8 +551,8 @@ bool UITest::WaitForBookmarkBarVisibilityChange(BrowserProxy* browser,
 }
 #endif  // defined(OS_WIN)
 
-GURL UITest::GetActiveTabURL() {
-  scoped_ptr<TabProxy> tab_proxy(GetActiveTab());
+GURL UITest::GetActiveTabURL(int window_index) {
+  scoped_ptr<TabProxy> tab_proxy(GetActiveTab(window_index));
   if (!tab_proxy.get())
     return GURL();
 
@@ -556,9 +562,9 @@ GURL UITest::GetActiveTabURL() {
   return url;
 }
 
-std::wstring UITest::GetActiveTabTitle() {
+std::wstring UITest::GetActiveTabTitle(int window_index) {
   std::wstring title;
-  scoped_ptr<TabProxy> tab_proxy(GetActiveTab());
+  scoped_ptr<TabProxy> tab_proxy(GetActiveTab(window_index));
   if (!tab_proxy.get())
     return title;
 
@@ -566,8 +572,8 @@ std::wstring UITest::GetActiveTabTitle() {
   return title;
 }
 
-int UITest::GetActiveTabIndex() {
-  scoped_ptr<TabProxy> tab_proxy(GetActiveTab());
+int UITest::GetActiveTabIndex(int window_index) {
+  scoped_ptr<TabProxy> tab_proxy(GetActiveTab(window_index));
   if (!tab_proxy.get())
     return -1;
 
diff --git a/chrome/test/ui/ui_test.h b/chrome/test/ui/ui_test.h
index e3a7c79..ff9a87b 100644
--- a/chrome/test/ui/ui_test.h
+++ b/chrome/test/ui/ui_test.h
@@ -96,15 +96,28 @@ class UITest : public testing::Test {
   // This method doesn't return until the navigation is complete.
   void NavigateToURL(const GURL& url);
 
-  // Returns the URL of the currently active tab. If there is no active tab,
-  // or some other error, the returned URL will be empty.
-  GURL GetActiveTabURL();
+  // Returns the URL of the currently active tab. Only looks in the first
+  // window, for backward compatibility. If there is no active tab, or some
+  // other error, the returned URL will be empty.
+  GURL GetActiveTabURL() { return GetActiveTabURL(0); }
 
-  // Returns the title of the currently active tab.
-  std::wstring GetActiveTabTitle();
+  // Like above, but looks at the window at the given index.
+  GURL GetActiveTabURL(int window_index);
 
-  // Returns the tabstrip index of the currently active tab, or -1 on error.
-  int GetActiveTabIndex();
+  // Returns the title of the currently active tab. Only looks in the first
+  // window, for backward compatibility.
+  std::wstring GetActiveTabTitle() { return GetActiveTabTitle(0); }
+
+  // Like above, but looks at the window at the given index.
+  std::wstring GetActiveTabTitle(int window_index);
+
+  // Returns the tabstrip index of the currently active tab in the window at
+  // the given index, or -1 on error. Only looks in the first window, for
+  // backward compatibility.
+  int GetActiveTabIndex() { return GetActiveTabIndex(0); }
+
+  // Like above, but looks at the window at the given index.
+  int GetActiveTabIndex(int window_index);
 
   // Returns true when the browser process is running, independent if any
   // renderer process exists or not. It will returns false if an user closed the
@@ -410,9 +423,13 @@ class UITest : public testing::Test {
   void CleanupAppProcesses();
 
   // Returns the proxy for the currently active tab, or NULL if there is no
-  // tab or there was some kind of error. The returned pointer MUST be
-  // deleted by the caller if non-NULL.
-  TabProxy* GetActiveTab();
+  // tab or there was some kind of error. Only looks at the first window, for
+  // backward compatibility. The returned pointer MUST be deleted by the
+  // caller if non-NULL.
+  TabProxy* GetActiveTab() { return GetActiveTab(0); }
+
+  // Like above, but looks at the window at the given index.
+  TabProxy* GetActiveTab(int window_index);
 
   // ********* Member variables *********