[Read aloud] Remove setSimpleNodeStoreWithText

Bug: 440400392
Change-Id: If6799edbf3c4b71985e090c98c3894ed2dbf6377
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/6991052
Auto-Submit: Kristi Saney <kristislee@google.com>
Commit-Queue: Lauren Winston <lwinston@google.com>
Reviewed-by: Lauren Winston <lwinston@google.com>
Reviewed-by: Eitan Goldberger <eitang@google.com>
Cr-Commit-Position: refs/heads/main@{#1522462}
diff --git a/chrome/test/data/webui/side_panel/read_anything/app_content_test.ts b/chrome/test/data/webui/side_panel/read_anything/app_content_test.ts
index e5a8429..e1ab103 100644
--- a/chrome/test/data/webui/side_panel/read_anything/app_content_test.ts
+++ b/chrome/test/data/webui/side_panel/read_anything/app_content_test.ts
@@ -8,7 +8,7 @@
 import {assertEquals, assertFalse, assertStringContains, assertTrue} from 'chrome-untrusted://webui-test/chai_assert.js';
 import {microtasksFinished} from 'chrome-untrusted://webui-test/test_util.js';
 
-import {createApp, emitEvent, setSimpleNodeStoreWithTextAndModel, setSimpleTreeWithText, setupBasicSpeech} from './common.js';
+import {createApp, emitEvent, setContent, setupBasicSpeech} from './common.js';
 import {FakeReadingMode} from './fake_reading_mode.js';
 import {TestColorUpdaterBrowserProxy} from './test_color_updater_browser_proxy.js';
 import {TestReadAloudModelBrowserProxy} from './test_read_aloud_browser_proxy.js';
@@ -75,7 +75,7 @@
   });
 
   test('showLoading clears read aloud state', () => {
-    setSimpleTreeWithText('My name is Regina George');
+    setContent('My name is Regina George', readAloudModel);
     emitEvent(app, ToolbarEvent.PLAY_PAUSE);
     assertTrue(speechController.isSpeechActive());
 
@@ -448,7 +448,7 @@
     });
 
     test('selection disallowed when speech active', async () => {
-      setSimpleNodeStoreWithTextAndModel('Been there, done that');
+      setContent('Been there, done that', readAloudModel);
 
       emitEvent(app, ToolbarEvent.PLAY_PAUSE);
       await microtasksFinished();
@@ -460,7 +460,7 @@
     });
 
     test('selection allowed after speech stops', async () => {
-      setSimpleNodeStoreWithTextAndModel('Who do you think you\'re kidding?');
+      setContent('Who do you think you\'re kidding?', readAloudModel);
 
       emitEvent(app, ToolbarEvent.PLAY_PAUSE);
       await microtasksFinished();
diff --git a/chrome/test/data/webui/side_panel/read_anything/app_receives_toolbar_changes_test.ts b/chrome/test/data/webui/side_panel/read_anything/app_receives_toolbar_changes_test.ts
index da89d4b..505cd71f 100644
--- a/chrome/test/data/webui/side_panel/read_anything/app_receives_toolbar_changes_test.ts
+++ b/chrome/test/data/webui/side_panel/read_anything/app_receives_toolbar_changes_test.ts
@@ -5,14 +5,15 @@
 import 'chrome-untrusted://read-anything-side-panel.top-chrome/read_anything.js';
 
 import type {AppElement} from 'chrome-untrusted://read-anything-side-panel.top-chrome/read_anything.js';
-import {BrowserProxy, SpeechBrowserProxyImpl, SpeechController, ToolbarEvent, VoiceLanguageController} from 'chrome-untrusted://read-anything-side-panel.top-chrome/read_anything.js';
+import {BrowserProxy, setInstance, SpeechBrowserProxyImpl, SpeechController, ToolbarEvent, VoiceLanguageController} from 'chrome-untrusted://read-anything-side-panel.top-chrome/read_anything.js';
 import {assertArrayEquals, assertEquals, assertFalse, assertNotEquals, assertTrue} from 'chrome-untrusted://webui-test/chai_assert.js';
 import {hasStyle, microtasksFinished} from 'chrome-untrusted://webui-test/test_util.js';
 
-import {createApp, createSpeechSynthesisVoice, emitEvent, mockMetrics, setSimpleTreeWithText, setupBasicSpeech} from './common.js';
+import {createApp, createSpeechSynthesisVoice, emitEvent, mockMetrics, setContent, setupBasicSpeech} from './common.js';
 import {FakeReadingMode} from './fake_reading_mode.js';
 import {TestColorUpdaterBrowserProxy} from './test_color_updater_browser_proxy.js';
 import type {TestMetricsBrowserProxy} from './test_metrics_browser_proxy.js';
+import {TestReadAloudModelBrowserProxy} from './test_read_aloud_browser_proxy.js';
 import {TestSpeechBrowserProxy} from './test_speech_browser_proxy.js';
 
 suite('AppReceivesToolbarChanges', () => {
@@ -21,6 +22,7 @@
   let metrics: TestMetricsBrowserProxy;
   let voiceLanguageController: VoiceLanguageController;
   let speechController: SpeechController;
+  let readAloudModel: TestReadAloudModelBrowserProxy;
 
   function containerLetterSpacing(): number {
     return +window.getComputedStyle(app.$.container)
@@ -89,6 +91,8 @@
     const readingMode = new FakeReadingMode();
     chrome.readingMode = readingMode as unknown as typeof chrome.readingMode;
     metrics = mockMetrics();
+    readAloudModel = new TestReadAloudModelBrowserProxy();
+    setInstance(readAloudModel);
     voiceLanguageController = new VoiceLanguageController();
     VoiceLanguageController.setInstance(voiceLanguageController);
     speechController = new SpeechController();
@@ -228,7 +232,8 @@
 
   test('on speech rate change speech rate updated', async () => {
     setupBasicSpeech(speech);
-    setSimpleTreeWithText('we mean no harm');
+    readAloudModel.setInitialized(true);
+    setContent('we mean no harm', readAloudModel);
     app.updateContent();
     await emitPlayPause();
 
@@ -260,12 +265,11 @@
 
   suite('play/pause', () => {
     setup(() => {
-      app.updateContent();
-      return microtasksFinished();
+      readAloudModel.setInitialized(true);
+      setContent('We come in peace', readAloudModel);
     });
 
     test('on first click starts speech', async () => {
-      setSimpleTreeWithText('We come in peace');
       await emitPlayPause();
       assertTrue(speechController.isSpeechActive());
       assertTrue(speechController.isSpeechTreeInitialized());
@@ -273,7 +277,6 @@
     });
 
     test('on second click stops speech', async () => {
-      setSimpleTreeWithText('Don\'t be alarmed!');
       await emitPlayPause();
       await emitPlayPause();
 
@@ -343,7 +346,8 @@
   suite('on granularity change', () => {
     setup(() => {
       setupBasicSpeech(speech);
-      setSimpleTreeWithText('we mean no harm');
+      readAloudModel.setInitialized(true);
+      setContent('we mean no harm', readAloudModel);
       app.updateContent();
       return emitPlayPause();
     });
diff --git a/chrome/test/data/webui/side_panel/read_anything/common.ts b/chrome/test/data/webui/side_panel/read_anything/common.ts
index dddd499..8193e21 100644
--- a/chrome/test/data/webui/side_panel/read_anything/common.ts
+++ b/chrome/test/data/webui/side_panel/read_anything/common.ts
@@ -129,36 +129,8 @@
       overrides || {});
 }
 
-
-export function setSimpleAxTreeWithText(text: string) {
-  const axTree = {
-    rootId: 1,
-    nodes: [
-      {
-        id: 1,
-        role: 'rootWebArea',
-        htmlTag: '#document',
-        childIds: [2],
-      },
-      {id: 2, role: 'staticText', name: text},
-    ],
-  };
-  chrome.readingMode.setContentForTesting(axTree, [2]);
-}
-
-// TODO: crbug.com/440400392- Merge setSimpleNodeStoreWithText with
-// setSimpleNodeStoreWithTextAndModel
-function setSimpleNodeStoreWithText(text: string) {
-  const id = 2;
-  chrome.readingMode.getCurrentTextSegments =
-      () => [{nodeId: id, start: 0, length: text.length}];
-  chrome.readingMode.getCurrentTextContent = () => text;
-  chrome.readingMode.getTextContent = () => text;
-  NodeStore.getInstance().setDomNode(document.createTextNode(text), id);
-}
-
-export function setSimpleNodeStoreWithTextAndModel(
-    text: string, model?: TestReadAloudModelBrowserProxy): Node {
+export function setContent(
+    text: string, model: TestReadAloudModelBrowserProxy): Node {
   const id = 2;
   const node = document.createTextNode(text);
   NodeStore.getInstance().setDomNode(node, id);
@@ -170,10 +142,3 @@
   }
   return node;
 }
-
-
-// Sets up the simple AX tree and node store with the given text.
-export function setSimpleTreeWithText(text: string) {
-  setSimpleAxTreeWithText(text);
-  setSimpleNodeStoreWithText(text);
-}
diff --git a/chrome/test/data/webui/side_panel/read_anything/speech_controller_content_test.ts b/chrome/test/data/webui/side_panel/read_anything/speech_controller_content_test.ts
index 58ac8e1..8447f47 100644
--- a/chrome/test/data/webui/side_panel/read_anything/speech_controller_content_test.ts
+++ b/chrome/test/data/webui/side_panel/read_anything/speech_controller_content_test.ts
@@ -6,7 +6,7 @@
 import {assertEquals, assertFalse, assertTrue} from 'chrome-untrusted://webui-test/chai_assert.js';
 import {MockTimer} from 'chrome-untrusted://webui-test/mock_timer.js';
 
-import {createApp, createSpeechSynthesisVoice, setSimpleNodeStoreWithTextAndModel, stubAnimationFrame} from './common.js';
+import {createApp, createSpeechSynthesisVoice, setContent, stubAnimationFrame} from './common.js';
 import {FakeReadingMode} from './fake_reading_mode.js';
 import {TestColorUpdaterBrowserProxy} from './test_color_updater_browser_proxy.js';
 import {TestReadAloudModelBrowserProxy} from './test_read_aloud_browser_proxy.js';
@@ -123,7 +123,7 @@
       // Before any content has been set, init is not called.
       assertEquals(0, readAloudModel.getCallCount('init'));
 
-      setSimpleNodeStoreWithTextAndModel('hello', readAloudModel);
+      setContent('hello', readAloudModel);
       app.updateContent();
       assertEquals(1, readAloudModel.getCallCount('init'));
 
@@ -137,11 +137,11 @@
       await createApp();
       assertEquals(1, readAloudModel.getCallCount('resetModel'));
 
-      setSimpleNodeStoreWithTextAndModel('hello', readAloudModel);
+      setContent('hello', readAloudModel);
       app.updateContent();
       assertEquals(2, readAloudModel.getCallCount('resetModel'));
 
-      setSimpleNodeStoreWithTextAndModel('hello, it\'s me', readAloudModel);
+      setContent('hello, it\'s me', readAloudModel);
       app.updateContent();
       assertEquals(3, readAloudModel.getCallCount('resetModel'));
     });
@@ -151,11 +151,11 @@
       await createApp();
       assertEquals(0, readAloudModel.getCallCount('resetModel'));
 
-      setSimpleNodeStoreWithTextAndModel('hello', readAloudModel);
+      setContent('hello', readAloudModel);
       app.updateContent();
       assertEquals(0, readAloudModel.getCallCount('resetModel'));
 
-      setSimpleNodeStoreWithTextAndModel('hello, it\'s me', readAloudModel);
+      setContent('hello, it\'s me', readAloudModel);
       app.updateContent();
       assertEquals(0, readAloudModel.getCallCount('resetModel'));
     });
@@ -175,7 +175,7 @@
 
   test('clearReadAloudState', () => {
     const text = 'And I am a massive deal';
-    const node: Node = setSimpleNodeStoreWithTextAndModel(text, readAloudModel);
+    const node: Node = setContent(text, readAloudModel);
     wordBoundaries.updateBoundary(4);
     chrome.readingMode.onHighlightGranularityChanged(
         chrome.readingMode.sentenceHighlighting);
@@ -211,7 +211,7 @@
     };
     const text = 'You bring the corsets';
     readAloudModel.setInitialized(true);
-    const node = setSimpleNodeStoreWithTextAndModel(text, readAloudModel);
+    const node = setContent(text, readAloudModel);
 
     speechController.onPlayPauseToggle(node as HTMLElement);
     const spoken = await speech.whenCalled('speak');
@@ -253,8 +253,7 @@
       'onPlayPauseToggle resume with word boundaries cancels and re-speaks',
       () => {
         const textContent = 'And our fame and our faces';
-        const node =
-            setSimpleNodeStoreWithTextAndModel(textContent, readAloudModel);
+        const node = setContent(textContent, readAloudModel);
         speechController.onPlayPauseToggle(node as HTMLElement);
         speechController.onPlayPauseToggle(node as HTMLElement);
         wordBoundaries.updateBoundary(10);
@@ -377,8 +376,7 @@
   });
 
   test('onNextGranularityClick updates state', () => {
-    setSimpleNodeStoreWithTextAndModel(
-        'Know all about the glories', readAloudModel);
+    setContent('Know all about the glories', readAloudModel);
     wordBoundaries.updateBoundary(5);
     assertEquals(1, speech.getCallCount('cancel'));
 
@@ -390,7 +388,7 @@
   });
 
   test('onPreviousGranularityClick updates state', () => {
-    setSimpleNodeStoreWithTextAndModel('And the disgraces', readAloudModel);
+    setContent('And the disgraces', readAloudModel);
     wordBoundaries.updateBoundary(5);
     assertEquals(1, speech.getCallCount('cancel'));
 
@@ -403,7 +401,7 @@
 
   test('onVoiceMenuClose resume speech only if it was active before', () => {
     const text = 'You must agree that baby';
-    setSimpleNodeStoreWithTextAndModel(text, readAloudModel);
+    setContent(text, readAloudModel);
     speechController.onVoiceMenuOpen();
 
     speechController.onVoiceMenuClose();
@@ -426,7 +424,7 @@
 
   test('set previous reading position without saved state does nothing', () => {
     const text = 'But I took your hand';
-    setSimpleNodeStoreWithTextAndModel(text, readAloudModel);
+    setContent(text, readAloudModel);
     wordBoundaries.updateBoundary(4);
     chrome.readingMode.onHighlightGranularityChanged(
         chrome.readingMode.sentenceHighlighting);
@@ -448,7 +446,7 @@
 
   test('set previous reading position restores saved state', () => {
     const text = 'And promised I\'d withstand';
-    setSimpleNodeStoreWithTextAndModel(text, readAloudModel);
+    setContent(text, readAloudModel);
     wordBoundaries.updateBoundary(4);
     chrome.readingMode.onHighlightGranularityChanged(
         chrome.readingMode.sentenceHighlighting);
@@ -472,7 +470,7 @@
   test('onTabMuteStateChange updates speech volume', async () => {
     const text = 'We\'ll bring the cinches';
     readAloudModel.setInitialized(true);
-    setSimpleNodeStoreWithTextAndModel(text, readAloudModel);
+    setContent(text, readAloudModel);
 
     speechController.onTabMuteStateChange(true);
     onPlayPauseToggle(text);
diff --git a/chrome/test/data/webui/side_panel/read_anything/speech_controller_test.ts b/chrome/test/data/webui/side_panel/read_anything/speech_controller_test.ts
index 7d03217..46cd0549 100644
--- a/chrome/test/data/webui/side_panel/read_anything/speech_controller_test.ts
+++ b/chrome/test/data/webui/side_panel/read_anything/speech_controller_test.ts
@@ -1,13 +1,14 @@
 // Copyright 2025 The Chromium Authors
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
-import {BrowserProxy, MAX_SPEECH_LENGTH, NodeStore, ReadAloudHighlighter, SelectionController, SpeechBrowserProxyImpl, SpeechController, VoiceLanguageController, WordBoundaries} from 'chrome-untrusted://read-anything-side-panel.top-chrome/read_anything.js';
-import {assertEquals, assertFalse, assertGT, assertNotEquals, assertTrue} from 'chrome-untrusted://webui-test/chai_assert.js';
+import {BrowserProxy, MAX_SPEECH_LENGTH, NodeStore, ReadAloudHighlighter, SelectionController, setInstance, SpeechBrowserProxyImpl, SpeechController, VoiceLanguageController, WordBoundaries} from 'chrome-untrusted://read-anything-side-panel.top-chrome/read_anything.js';
+import {assertEquals, assertFalse, assertGE, assertGT, assertNotEquals, assertTrue} from 'chrome-untrusted://webui-test/chai_assert.js';
 
-import {createSpeechErrorEvent, createSpeechSynthesisVoice, createWordBoundaryEvent, mockMetrics, setSimpleTreeWithText} from './common.js';
+import {createSpeechErrorEvent, createSpeechSynthesisVoice, createWordBoundaryEvent, mockMetrics, setContent} from './common.js';
 import {FakeReadingMode} from './fake_reading_mode.js';
 import {TestColorUpdaterBrowserProxy} from './test_color_updater_browser_proxy.js';
 import type {TestMetricsBrowserProxy} from './test_metrics_browser_proxy.js';
+import {TestReadAloudModelBrowserProxy} from './test_read_aloud_browser_proxy.js';
 import {TestSpeechBrowserProxy} from './test_speech_browser_proxy.js';
 
 suite('SpeechController', () => {
@@ -24,10 +25,12 @@
   let voiceLanguageController: VoiceLanguageController;
   let readingMode: FakeReadingMode;
   let selectionController: SelectionController;
+  let readAloudModel: TestReadAloudModelBrowserProxy;
 
   // TODO: crbug.com/440400392- Move all tests relying on chrome.readingMode
   // for text segmentation to use TestReadAloudModelBrowserProxy instead.
   function onPlayPauseToggle(text: string) {
+    setContent(text, readAloudModel);
     const element = document.createElement('p');
     element.textContent = text;
     speechController.onPlayPauseToggle(element);
@@ -68,6 +71,9 @@
       },
     };
 
+    readAloudModel = new TestReadAloudModelBrowserProxy();
+    setInstance(readAloudModel);
+    readAloudModel.setInitialized(true);
     voiceLanguageController = new VoiceLanguageController();
     voiceLanguageController.setUserPreferredVoice(
         createSpeechSynthesisVoice({lang: 'en', name: 'Google Alpaca'}));
@@ -154,7 +160,7 @@
 
   test('onSpeechSettingsChange cancels and resumes speech if playing', () => {
     const text = 'In all the time I\'ve been by your side';
-    setSimpleTreeWithText(text);
+    setContent(text, readAloudModel);
     onPlayPauseToggle(text);
     speech.reset();
 
@@ -171,7 +177,7 @@
 
   test('onSpeechSettingsChange does not resume speech if not playing', () => {
     speechController.setHasSpeechBeenTriggered(true);
-    setSimpleTreeWithText('I\'ve never lost control');
+    setContent('I\'ve never lost control', readAloudModel);
 
     speechController.onSpeechSettingsChange();
 
@@ -198,7 +204,7 @@
 
   test('onPlayPauseToggle waits for engine load', async () => {
     const text = 'Sorry not sorry bout what I said';
-    setSimpleTreeWithText(text);
+    setContent(text, readAloudModel);
 
     onPlayPauseToggle(text);
     const spoken = await speech.whenCalled('speak');
@@ -222,7 +228,7 @@
     const text = 'I\'m just tryna have some fun';
     chrome.readingMode.speechRate = rate;
     chrome.readingMode.baseLanguageForSpeech = lang;
-    setSimpleTreeWithText(text);
+    setContent(text, readAloudModel);
 
     onPlayPauseToggle(text);
 
@@ -233,6 +239,7 @@
 
   test('onPlayPauseToggle pauses with button click', () => {
     onPlayPauseToggle('A story that you think');
+    speech.reset();
     onPlayPauseToggle('A story that you think');
 
     assertTrue(isSpeechActiveChanged);
@@ -265,7 +272,7 @@
 
   test('word boundary received updates words heard', () => {
     const textContent = 'You\'re all I can think of';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     onPlayPauseToggle(textContent);
     const spoken = speech.getArgs('speak')[0];
 
@@ -277,7 +284,7 @@
 
   test('words heard not updated for whitespace', () => {
     const textContent = 'Every drop I drink up';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     onPlayPauseToggle(textContent);
     const spoken = speech.getArgs('speak')[0];
 
@@ -289,7 +296,7 @@
 
   test('words heard reset on clear', () => {
     const textContent = 'You\'re my soda pop';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     onPlayPauseToggle(textContent);
     const spoken = speech.getArgs('speak')[0];
 
@@ -304,7 +311,7 @@
 
   test('sentence end with word boundaries, does not count sentence', () => {
     const textContent = 'My little soda pop';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     onPlayPauseToggle(textContent);
     const spoken = speech.getArgs('speak')[0];
 
@@ -317,7 +324,7 @@
 
   test('sentence end with no word boundaries, counts sentence', () => {
     const textContent = 'Cool me down, you\'re so hot';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     onPlayPauseToggle(textContent);
     const spoken = speech.getArgs('speak')[0];
 
@@ -350,7 +357,7 @@
         'here I stay- let the storm rage on';
 
     setup(() => {
-      setSimpleTreeWithText(longSentences);
+      setContent(longSentences, readAloudModel);
     });
 
     test('uses max speech length', () => {
@@ -362,7 +369,7 @@
       assertGT(expectedNumSegments, 0);
       for (let i = 0; i < expectedNumSegments; i++) {
         assertEquals(i + 1, speech.getCallCount('speak'));
-        assertGT(
+        assertGE(
             MAX_SPEECH_LENGTH, speech.getArgs('speak')[i].text.trim().length);
         speech.getArgs('speak')[i].onend();
       }
@@ -394,7 +401,7 @@
   test('stops speech on language-unavailable', async () => {
     const textContent = 'I\'m done cuz all this time';
     const pageLanguage = 'es';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     assertNotEquals(chrome.readingMode.defaultLanguageForSpeech, pageLanguage);
     chrome.readingMode.baseLanguageForSpeech = pageLanguage;
     voiceLanguageController.onPageLanguageChanged();
@@ -420,7 +427,7 @@
   test('stops speech on voice-unavailable', async () => {
     const textContent = 'I\'ve been just one word';
     const pageLanguage = 'es';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     assertNotEquals(chrome.readingMode.defaultLanguageForSpeech, pageLanguage);
     chrome.readingMode.baseLanguageForSpeech = pageLanguage;
     voiceLanguageController.onPageLanguageChanged();
@@ -445,7 +452,7 @@
   test('invalid argument updates speech rate', () => {
     const textContent = 'In a stupid rhyme';
     const pageLanguage = 'es';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     assertNotEquals(chrome.readingMode.defaultLanguageForSpeech, pageLanguage);
     chrome.readingMode.speechRate = 4;
     chrome.readingMode.baseLanguageForSpeech = pageLanguage;
@@ -470,7 +477,7 @@
   test('speech interrupt while repositioning keeps playing speech', () => {
     const textContent = 'So I picked up a pen and a microphone';
     const pageLanguage = 'es';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     assertNotEquals(chrome.readingMode.defaultLanguageForSpeech, pageLanguage);
     chrome.readingMode.speechRate = 4;
     chrome.readingMode.baseLanguageForSpeech = pageLanguage;
@@ -497,7 +504,7 @@
   test('speech interrupt stops speech', async () => {
     const textContent = 'History\'s about to get overthrown';
     const pageLanguage = 'es';
-    setSimpleTreeWithText(textContent);
+    setContent(textContent, readAloudModel);
     assertNotEquals(chrome.readingMode.defaultLanguageForSpeech, pageLanguage);
     chrome.readingMode.speechRate = 4;
     chrome.readingMode.baseLanguageForSpeech = pageLanguage;
@@ -528,11 +535,7 @@
 
   test('speech finished clears state', async () => {
     const text = 'New phone who dis?';
-    setSimpleTreeWithText(text);
-    let resetGranularityIndex = false;
-    chrome.readingMode.resetGranularityIndex = () => {
-      resetGranularityIndex = true;
-    };
+    setContent(text, readAloudModel);
 
     onPlayPauseToggle(text);
 
@@ -541,11 +544,11 @@
 
     speech.reset();
     isSpeechActiveChanged = false;
-    chrome.readingMode.getCurrentTextSegments = () => [];
+    readAloudModel.setCurrentTextSegments([]);
     spoken.onend();
 
     assertTrue(isSpeechActiveChanged);
-    assertTrue(resetGranularityIndex);
+    assertEquals(1, readAloudModel.getCallCount('resetSpeechToBeginning'));
     assertFalse(speechController.isSpeechActive());
     assertFalse(speechController.isPausedFromButton());
     assertFalse(speechController.isTemporaryPause());
@@ -556,32 +559,18 @@
   });
 
   test('onNextGranularityClick propagates change', () => {
-    let movedToNext = false;
-    chrome.readingMode.getCurrentTextSegments = () => [];
-    chrome.readingMode.movePositionToNextGranularity = () => {
-      movedToNext = true;
-    };
-
     speechController.onNextGranularityClick();
-
-    assertTrue(movedToNext);
+    assertEquals(1, readAloudModel.getCallCount('moveSpeechForward'));
   });
 
   test('onPreviousGranularityClick propagates change', () => {
-    let movedToPrevious: boolean = false;
-    chrome.readingMode.getCurrentTextSegments = () => [];
-    chrome.readingMode.movePositionToPreviousGranularity = () => {
-      movedToPrevious = true;
-    };
-
     speechController.onPreviousGranularityClick();
-
-    assertTrue(movedToPrevious);
+    assertEquals(1, readAloudModel.getCallCount('moveSpeechBackwards'));
   });
 
   test('onHighlightGranularityChange draws highlight', () => {
     const granularity = chrome.readingMode.wordHighlighting;
-    setSimpleTreeWithText('no more melon cake');
+    setContent('no more melon cake', readAloudModel);
     assertFalse(highlighter.hasCurrentGranularity());
 
     speechController.onHighlightGranularityChange(granularity);