diff --git a/test/test.js b/test/test.js
index 5baee43d9a..402366f2d7 100644
--- a/test/test.js
+++ b/test/test.js
@@ -17,6 +17,9 @@ const { By, until } = webdriver;
 const path = require("path");
 
 const SHOOTER_BUTTON_ID = "pageAction-panel-screenshots";
+// Applies to the old-style toolbar button:
+const TOOLBAR_SHOOTER_BUTTON_ID = "screenshots_mozilla_org-browser-action";
+const shooterSelector = By.css(`#${SHOOTER_BUTTON_ID}, #${TOOLBAR_SHOOTER_BUTTON_ID}`);
 const SLIDE_IFRAME_ID = "firefox-screenshots-onboarding-iframe";
 const PRESELECTION_IFRAME_ID = "firefox-screenshots-preselection-iframe";
 const PREVIEW_IFRAME_ID = "firefox-screenshots-preview-iframe";
@@ -118,7 +121,7 @@ function promiseFinally(promise, finallyCallback) {
     successfully. */
 function startScreenshots(driver) {
   return promiseFinally(
-    getChromeElement(driver, By.id(SHOOTER_BUTTON_ID)).then((button) => {
+    getChromeElement(driver, shooterSelector).then((button) => {
       return button.click();
     }),
     () => {
@@ -219,10 +222,14 @@ describe("Test Screenshots", function() {
   it("should find the add-on button", function() {
     this.timeout(15000);
     return promiseFinally(
-      getChromeElement(driver, By.id(SHOOTER_BUTTON_ID))
+      getChromeElement(driver, shooterSelector)
       .then((button) => button.getAttribute("label"))
       .then((label) => {
-        assert.equal(label, "Take a Screenshot");
+        if (label == "Take a Screenshot") {
+          assert.equal(label, "Take a Screenshot");
+        } else {
+          assert.equal(label, "Firefox Screenshots");
+        }
       }),
       () => {
         driver.setContext(firefox.Context.CONTENT);