4 files changed, 351 insertions, 128 deletions

diff --git a/screen_ui.cpp b/screen_ui.cpp
index 391dedb00..181d58ecf 100644
--- a/screen_ui.cpp
+++ b/screen_ui.cpp
@@ -54,15 +54,23 @@ static double now() {
   return tv.tv_sec + tv.tv_usec / 1000000.0;
 }
 
-Menu::Menu(bool scrollable, size_t max_items, size_t max_length,
-           const std::vector<std::string>& headers, const std::vector<std::string>& items,
-           size_t initial_selection)
-    : scrollable_(scrollable),
+Menu::Menu(size_t initial_selection, const DrawInterface& draw_func)
+    : selection_(initial_selection), draw_funcs_(draw_func) {}
+
+size_t Menu::selection() const {
+  return selection_;
+}
+
+TextMenu::TextMenu(bool scrollable, size_t max_items, size_t max_length,
+                   const std::vector<std::string>& headers, const std::vector<std::string>& items,
+                   size_t initial_selection, int char_height, const DrawInterface& draw_funcs)
+    : Menu(initial_selection, draw_funcs),
+      scrollable_(scrollable),
       max_display_items_(max_items),
       max_item_length_(max_length),
       text_headers_(headers),
       menu_start_(0),
-      selection_(initial_selection) {
+      char_height_(char_height) {
   CHECK_LE(max_items, static_cast<size_t>(std::numeric_limits<int>::max()));
 
   // It's fine to have more entries than text_rows_ if scrollable menu is supported.
@@ -74,29 +82,29 @@ Menu::Menu(bool scrollable, size_t max_items, size_t max_length,
   CHECK(!text_items_.empty());
 }
 
-const std::vector<std::string>& Menu::text_headers() const {
+const std::vector<std::string>& TextMenu::text_headers() const {
   return text_headers_;
 }
 
-std::string Menu::TextItem(size_t index) const {
+std::string TextMenu::TextItem(size_t index) const {
   CHECK_LT(index, text_items_.size());
 
   return text_items_[index];
 }
 
-size_t Menu::MenuStart() const {
+size_t TextMenu::MenuStart() const {
   return menu_start_;
 }
 
-size_t Menu::MenuEnd() const {
+size_t TextMenu::MenuEnd() const {
   return std::min(ItemsCount(), menu_start_ + max_display_items_);
 }
 
-size_t Menu::ItemsCount() const {
+size_t TextMenu::ItemsCount() const {
   return text_items_.size();
 }
 
-bool Menu::ItemsOverflow(std::string* cur_selection_str) const {
+bool TextMenu::ItemsOverflow(std::string* cur_selection_str) const {
   if (!scrollable_ || ItemsCount() <= max_display_items_) {
     return false;
   }
@@ -107,7 +115,7 @@ bool Menu::ItemsOverflow(std::string* cur_selection_str) const {
 }
 
 // TODO(xunchang) modify the function parameters to button up & down.
-int Menu::Select(int sel) {
+int TextMenu::Select(int sel) {
   CHECK_LE(ItemsCount(), static_cast<size_t>(std::numeric_limits<int>::max()));
   int count = ItemsCount();
 
@@ -140,6 +148,151 @@ int Menu::Select(int sel) {
   return selection_;
 }
 
+int TextMenu::DrawHeader(int x, int y) const {
+  int offset = 0;
+
+  draw_funcs_.SetColor(UIElement::HEADER);
+  if (!scrollable()) {
+    offset += draw_funcs_.DrawWrappedTextLines(x, y + offset, text_headers());
+  } else {
+    offset += draw_funcs_.DrawTextLines(x, y + offset, text_headers());
+    // Show the current menu item number in relation to total number if items don't fit on the
+    // screen.
+    std::string cur_selection_str;
+    if (ItemsOverflow(&cur_selection_str)) {
+      offset += draw_funcs_.DrawTextLine(x, y + offset, cur_selection_str, true);
+    }
+  }
+
+  return offset;
+}
+
+int TextMenu::DrawItems(int x, int y, int screen_width, bool long_press) const {
+  int offset = 0;
+
+  draw_funcs_.SetColor(UIElement::MENU);
+  // Do not draw the horizontal rule for wear devices.
+  if (!scrollable()) {
+    offset += draw_funcs_.DrawHorizontalRule(y + offset) + 4;
+  }
+  for (size_t i = MenuStart(); i < MenuEnd(); ++i) {
+    bool bold = false;
+    if (i == selection()) {
+      // Draw the highlight bar.
+      draw_funcs_.SetColor(long_press ? UIElement::MENU_SEL_BG_ACTIVE : UIElement::MENU_SEL_BG);
+
+      int bar_height = char_height_ + 4;
+      draw_funcs_.DrawHighlightBar(0, y + offset - 2, screen_width, bar_height);
+
+      // Bold white text for the selected item.
+      draw_funcs_.SetColor(UIElement::MENU_SEL_FG);
+      bold = true;
+    }
+    offset += draw_funcs_.DrawTextLine(x, y + offset, TextItem(i), bold);
+
+    draw_funcs_.SetColor(UIElement::MENU);
+  }
+  offset += draw_funcs_.DrawHorizontalRule(y + offset);
+
+  return offset;
+}
+
+GraphicMenu::GraphicMenu(size_t max_width, size_t max_height, GRSurface* graphic_headers,
+                         const std::vector<GRSurface*>& graphic_items, size_t initial_selection,
+                         const DrawInterface& draw_funcs)
+    : Menu(initial_selection, draw_funcs),
+      max_width_(max_width),
+      max_height_(max_height),
+      graphic_headers_(graphic_headers),
+      graphic_items_(graphic_items) {}
+
+int GraphicMenu::Select(int sel) {
+  CHECK_LE(graphic_items_.size(), static_cast<size_t>(std::numeric_limits<int>::max()));
+  int count = graphic_items_.size();
+
+  // Wraps the selection at boundary if the menu is not scrollable.
+  if (sel < 0) {
+    selection_ = count - 1;
+  } else if (sel >= count) {
+    selection_ = 0;
+  } else {
+    selection_ = sel;
+  }
+
+  return selection_;
+}
+
+int GraphicMenu::DrawHeader(int x, int y) const {
+  draw_funcs_.DrawTextIcon(x, y, graphic_headers_);
+  return graphic_headers_->height;
+}
+
+int GraphicMenu::DrawItems(int x, int y, int screen_width, bool long_press) const {
+  int offset = 0;
+
+  draw_funcs_.SetColor(UIElement::MENU);
+  offset += draw_funcs_.DrawHorizontalRule(y + offset) + 4;
+
+  for (size_t i = 0; i < graphic_items_.size(); i++) {
+    auto& item = graphic_items_[i];
+    if (i == selection_) {
+      draw_funcs_.SetColor(long_press ? UIElement::MENU_SEL_BG_ACTIVE : UIElement::MENU_SEL_BG);
+
+      int bar_height = item->height + 4;
+      draw_funcs_.DrawHighlightBar(0, y + offset - 2, screen_width, bar_height);
+
+      // Bold white text for the selected item.
+      draw_funcs_.SetColor(UIElement::MENU_SEL_FG);
+    }
+    draw_funcs_.DrawTextIcon(x, y + offset, item);
+    offset += item->height;
+
+    draw_funcs_.SetColor(UIElement::MENU);
+  }
+
+  return offset;
+}
+
+bool GraphicMenu::Validate() const {
+  int offset = 0;
+  if (!ValidateGraphicSurface(offset, graphic_headers_)) {
+    return false;
+  }
+  offset += graphic_headers_->height;
+
+  for (const auto& item : graphic_items_) {
+    if (!ValidateGraphicSurface(offset, item)) {
+      return false;
+    }
+    offset += item->height;
+  }
+
+  return true;
+}
+
+bool GraphicMenu::ValidateGraphicSurface(int y, const GRSurface* surface) const {
+  if (!surface) {
+    fprintf(stderr, "Graphic surface can not be null");
+    return false;
+  }
+
+  if (surface->pixel_bytes != 1 || surface->width != surface->row_bytes) {
+    fprintf(stderr, "Invalid graphic surface, pixel bytes: %d, width: %d row_bytes: %d",
+            surface->pixel_bytes, surface->width, surface->row_bytes);
+    return false;
+  }
+
+  if (surface->width > max_width_ || surface->height > max_height_ - y) {
+    fprintf(stderr,
+            "Graphic surface doesn't fit into the screen. width: %d, height: %d, max_width: %zu,"
+            " max_height: %zu, vertical offset: %d\n",
+            surface->width, surface->height, max_width_, max_height_, y);
+    return false;
+  }
+
+  return true;
+}
+
 ScreenRecoveryUI::ScreenRecoveryUI() : ScreenRecoveryUI(false) {}
 
 constexpr int kDefaultMarginHeight = 0;
@@ -332,26 +485,26 @@ void ScreenRecoveryUI::draw_foreground_locked() {
 
 void ScreenRecoveryUI::SetColor(UIElement e) const {
   switch (e) {
-    case INFO:
+    case UIElement::INFO:
       gr_color(249, 194, 0, 255);
       break;
-    case HEADER:
+    case UIElement::HEADER:
       gr_color(247, 0, 6, 255);
       break;
-    case MENU:
-    case MENU_SEL_BG:
+    case UIElement::MENU:
+    case UIElement::MENU_SEL_BG:
       gr_color(0, 106, 157, 255);
       break;
-    case MENU_SEL_BG_ACTIVE:
+    case UIElement::MENU_SEL_BG_ACTIVE:
       gr_color(0, 156, 100, 255);
       break;
-    case MENU_SEL_FG:
+    case UIElement::MENU_SEL_FG:
       gr_color(255, 255, 255, 255);
       break;
-    case LOG:
+    case UIElement::LOG:
       gr_color(196, 196, 196, 255);
       break;
-    case TEXT_FILL:
+    case UIElement::TEXT_FILL:
       gr_color(0, 0, 0, 160);
       break;
     default:
@@ -384,7 +537,7 @@ void ScreenRecoveryUI::SelectAndShowBackgroundText(const std::vector<std::string
   int text_x = margin_width_;
   int line_spacing = gr_sys_font()->char_height;  // Put some extra space between images.
   // Write the header and descriptive texts.
-  SetColor(INFO);
+  SetColor(UIElement::INFO);
   std::string header = "Show background text image";
   text_y += DrawTextLine(text_x, text_y, header, true);
   std::string locale_selection = android::base::StringPrintf(
@@ -400,7 +553,7 @@ void ScreenRecoveryUI::SelectAndShowBackgroundText(const std::vector<std::string
   // Iterate through the text images and display them in order for the current locale.
   for (const auto& p : surfaces) {
     text_y += line_spacing;
-    SetColor(LOG);
+    SetColor(UIElement::LOG);
     text_y += DrawTextLine(text_x, text_y, p.first, false);
     gr_color(255, 255, 255, 255);
     gr_texticon(text_x, text_y, p.second.get());
@@ -547,7 +700,7 @@ void ScreenRecoveryUI::draw_menu_and_text_buffer_locked(
     static constexpr int kMenuIndent = 4;
     int x = margin_width_ + kMenuIndent;
 
-    SetColor(INFO);
+    SetColor(UIElement::INFO);
 
     for (size_t i = 0; i < title_lines_.size(); i++) {
       y += DrawTextLine(x, y, title_lines_[i], i == 0);
@@ -555,50 +708,13 @@ void ScreenRecoveryUI::draw_menu_and_text_buffer_locked(
 
     y += DrawTextLines(x, y, help_message);
 
-    // Draw menu header.
-    SetColor(HEADER);
-    if (!menu_->scrollable()) {
-      y += DrawWrappedTextLines(x, y, menu_->text_headers());
-    } else {
-      y += DrawTextLines(x, y, menu_->text_headers());
-      // Show the current menu item number in relation to total number if items don't fit on the
-      // screen.
-      std::string cur_selection_str;
-      if (menu_->ItemsOverflow(&cur_selection_str)) {
-        y += DrawTextLine(x, y, cur_selection_str, true);
-      }
-    }
-
-    // Draw menu items.
-    SetColor(MENU);
-    // Do not draw the horizontal rule for wear devices.
-    if (!menu_->scrollable()) {
-      y += DrawHorizontalRule(y) + 4;
-    }
-    for (size_t i = menu_->MenuStart(); i < menu_->MenuEnd(); ++i) {
-      bool bold = false;
-      if (i == static_cast<size_t>(menu_->selection())) {
-        // Draw the highlight bar.
-        SetColor(IsLongPress() ? MENU_SEL_BG_ACTIVE : MENU_SEL_BG);
-
-        int bar_height = char_height_ + 4;
-        DrawHighlightBar(0, y - 2, ScreenWidth(), bar_height);
-
-        // Bold white text for the selected item.
-        SetColor(MENU_SEL_FG);
-        bold = true;
-      }
-
-      y += DrawTextLine(x, y, menu_->TextItem(i), bold);
-
-      SetColor(MENU);
-    }
-    y += DrawHorizontalRule(y);
+    y += menu_->DrawHeader(x, y);
+    y += menu_->DrawItems(x, y, ScreenWidth(), IsLongPress());
   }
 
   // Display from the bottom up, until we hit the top of the screen, the bottom of the menu, or
   // we've displayed the entire text buffer.
-  SetColor(LOG);
+  SetColor(UIElement::LOG);
   int row = text_row_;
   size_t count = 0;
   for (int ty = ScreenHeight() - margin_height_ - char_height_; ty >= y && count < text_rows_;
@@ -992,8 +1108,8 @@ void ScreenRecoveryUI::StartMenu(const std::vector<std::string>& headers,
                                  const std::vector<std::string>& items, size_t initial_selection) {
   std::lock_guard<std::mutex> lg(updateMutex);
   if (text_rows_ > 0 && text_cols_ > 1) {
-    menu_ = std::make_unique<Menu>(scrollable_menu_, text_rows_, text_cols_ - 1, headers, items,
-                                   initial_selection);
+    menu_ = std::make_unique<TextMenu>(scrollable_menu_, text_rows_, text_cols_ - 1, headers, items,
+                                       initial_selection, char_height_, *this);
     update_screen_locked();
   }
 }
diff --git a/screen_ui.h b/screen_ui.h
index f08f4f4f3..915288793 100644
--- a/screen_ui.h
+++ b/screen_ui.h
@@ -31,23 +31,92 @@
 // From minui/minui.h.
 struct GRSurface;
 
-// This class maintains the menu selection and display of the screen ui.
+enum class UIElement {
+  HEADER,
+  MENU,
+  MENU_SEL_BG,
+  MENU_SEL_BG_ACTIVE,
+  MENU_SEL_FG,
+  LOG,
+  TEXT_FILL,
+  INFO
+};
+
+// Interface to draw the UI elements on the screen.
+class DrawInterface {
+ public:
+  virtual ~DrawInterface() = default;
+
+  // Sets the color to the predefined value for |element|.
+  virtual void SetColor(UIElement element) const = 0;
+
+  // Draws a highlight bar at (x, y) - (x + width, y + height).
+  virtual void DrawHighlightBar(int x, int y, int width, int height) const = 0;
+
+  // Draws a horizontal rule at Y. Returns the offset it should be moving along Y-axis.
+  virtual int DrawHorizontalRule(int y) const = 0;
+
+  // Draws a line of text. Returns the offset it should be moving along Y-axis.
+  virtual int DrawTextLine(int x, int y, const std::string& line, bool bold) const = 0;
+
+  // Draws surface portion (sx, sy, w, h) at screen location (dx, dy).
+  virtual void DrawSurface(GRSurface* surface, int sx, int sy, int w, int h, int dx,
+                           int dy) const = 0;
+
+  // Draws rectangle at (x, y) - (x + w, y + h).
+  virtual void DrawFill(int x, int y, int w, int h) const = 0;
+
+  // Draws given surface (surface->pixel_bytes = 1) as text at (x, y).
+  virtual void DrawTextIcon(int x, int y, GRSurface* surface) const = 0;
+
+  // Draws multiple text lines. Returns the offset it should be moving along Y-axis.
+  virtual int DrawTextLines(int x, int y, const std::vector<std::string>& lines) const = 0;
+
+  // Similar to DrawTextLines() to draw multiple text lines, but additionally wraps long lines. It
+  // keeps symmetrical margins of 'x' at each end of a line. Returns the offset it should be moving
+  // along Y-axis.
+  virtual int DrawWrappedTextLines(int x, int y, const std::vector<std::string>& lines) const = 0;
+};
+
+// Interface for classes that maintain the menu selection and display.
 class Menu {
  public:
+  virtual ~Menu() = default;
+  // Returns the current menu selection.
+  size_t selection() const;
+  // Sets the current selection to |sel|. Handle the overflow cases depending on if the menu is
+  // scrollable.
+  virtual int Select(int sel) = 0;
+  // Displays the menu headers on the screen at offset x, y
+  virtual int DrawHeader(int x, int y) const = 0;
+  // Iterates over the menu items and displays each of them at offset x, y.
+  virtual int DrawItems(int x, int y, int screen_width, bool long_press) const = 0;
+
+ protected:
+  Menu(size_t initial_selection, const DrawInterface& draw_func);
+  // Current menu selection.
+  size_t selection_;
+  // Reference to the class that implements all the draw functions.
+  const DrawInterface& draw_funcs_;
+};
+
+// This class uses strings as the menu header and items.
+class TextMenu : public Menu {
+ public:
   // Constructs a Menu instance with the given |headers|, |items| and properties. Sets the initial
   // selection to |initial_selection|.
-  Menu(bool scrollable, size_t max_items, size_t max_length,
-       const std::vector<std::string>& headers, const std::vector<std::string>& items,
-       size_t initial_selection);
+  TextMenu(bool scrollable, size_t max_items, size_t max_length,
+           const std::vector<std::string>& headers, const std::vector<std::string>& items,
+           size_t initial_selection, int char_height, const DrawInterface& draw_funcs);
+
+  int Select(int sel) override;
+  int DrawHeader(int x, int y) const override;
+  int DrawItems(int x, int y, int screen_width, bool long_press) const override;
 
   bool scrollable() const {
     return scrollable_;
   }
 
-  size_t selection() const {
-    return selection_;
-  }
-
   // Returns count of menu items.
   size_t ItemsCount() const;
 
@@ -75,10 +144,6 @@ class Menu {
   // |cur_selection_str| if the items exceed the screen limit.
   bool ItemsOverflow(std::string* cur_selection_str) const;
 
-  // Sets the current selection to |sel|. Handle the overflow cases depending on if the menu is
-  // scrollable.
-  int Select(int sel);
-
  private:
   // The menu is scrollable to display more items. Used on wear devices who have smaller screens.
   const bool scrollable_;
@@ -92,25 +157,45 @@ class Menu {
   std::vector<std::string> text_items_;
   // The first item to display on the screen.
   size_t menu_start_;
-  // Current menu selection.
-  size_t selection_;
+
+  // Height in pixels of each character.
+  int char_height_;
+};
+
+// This class uses GRSurfaces* as the menu header and items.
+class GraphicMenu : public Menu {
+ public:
+  // Constructs a Menu instance with the given |headers|, |items| and properties. Sets the initial
+  // selection to |initial_selection|.
+  GraphicMenu(size_t max_width, size_t max_height, GRSurface* graphic_headers,
+              const std::vector<GRSurface*>& graphic_items, size_t initial_selection,
+              const DrawInterface& draw_funcs);
+
+  int Select(int sel) override;
+  int DrawHeader(int x, int y) const override;
+  int DrawItems(int x, int y, int screen_width, bool long_press) const override;
+
+  // Checks if all the header and items are valid GRSurfaces; and that they can fit in the area
+  // defined by |max_width_| and |max_height_|.
+  bool Validate() const;
+
+ private:
+  // Returns true if |surface| fits on the screen with a vertical offset |y|.
+  bool ValidateGraphicSurface(int y, const GRSurface* surface) const;
+
+  const size_t max_width_;
+  const size_t max_height_;
+
+  // Pointers to the menu headers and items in graphic icons. This class does not have the ownership
+  // of the these objects.
+  GRSurface* graphic_headers_;
+  std::vector<GRSurface*> graphic_items_;
 };
 
 // Implementation of RecoveryUI appropriate for devices with a screen
 // (shows an icon + a progress bar, text logging, menu, etc.)
-class ScreenRecoveryUI : public RecoveryUI {
+class ScreenRecoveryUI : public RecoveryUI, public DrawInterface {
  public:
-  enum UIElement {
-    HEADER,
-    MENU,
-    MENU_SEL_BG,
-    MENU_SEL_BG_ACTIVE,
-    MENU_SEL_FG,
-    LOG,
-    TEXT_FILL,
-    INFO
-  };
-
   ScreenRecoveryUI();
   explicit ScreenRecoveryUI(bool scrollable_menu);
   ~ScreenRecoveryUI() override;
@@ -149,8 +234,6 @@ class ScreenRecoveryUI : public RecoveryUI {
 
   void Redraw();
 
-  void SetColor(UIElement e) const;
-
   // Checks the background text image, for debugging purpose. It iterates the locales embedded in
   // the on-device resource files and shows the localized text, for manual inspection.
   void CheckBackgroundTextImages();
@@ -212,24 +295,16 @@ class ScreenRecoveryUI : public RecoveryUI {
   // Returns pixel height of draw buffer.
   virtual int ScreenHeight() const;
 
-  // Draws a highlight bar at (x, y) - (x + width, y + height).
-  virtual void DrawHighlightBar(int x, int y, int width, int height) const;
-  // Draws a horizontal rule at Y. Returns the offset it should be moving along Y-axis.
-  virtual int DrawHorizontalRule(int y) const;
-  // Draws a line of text. Returns the offset it should be moving along Y-axis.
-  virtual int DrawTextLine(int x, int y, const std::string& line, bool bold) const;
-  // Draws surface portion (sx, sy, w, h) at screen location (dx, dy).
-  virtual void DrawSurface(GRSurface* surface, int sx, int sy, int w, int h, int dx, int dy) const;
-  // Draws rectangle at (x, y) - (x + w, y + h).
-  virtual void DrawFill(int x, int y, int w, int h) const;
-  // Draws given surface (surface->pixel_bytes = 1) as text at (x, y).
-  virtual void DrawTextIcon(int x, int y, GRSurface* surface) const;
-  // Draws multiple text lines. Returns the offset it should be moving along Y-axis.
-  int DrawTextLines(int x, int y, const std::vector<std::string>& lines) const;
-  // Similar to DrawTextLines() to draw multiple text lines, but additionally wraps long lines. It
-  // keeps symmetrical margins of 'x' at each end of a line. Returns the offset it should be moving
-  // along Y-axis.
-  int DrawWrappedTextLines(int x, int y, const std::vector<std::string>& lines) const;
+  // Implementation of the draw functions in DrawInterface.
+  void SetColor(UIElement e) const override;
+  void DrawHighlightBar(int x, int y, int width, int height) const override;
+  int DrawHorizontalRule(int y) const override;
+  void DrawSurface(GRSurface* surface, int sx, int sy, int w, int h, int dx, int dy) const override;
+  void DrawFill(int x, int y, int w, int h) const override;
+  void DrawTextIcon(int x, int y, GRSurface* surface) const override;
+  int DrawTextLine(int x, int y, const std::string& line, bool bold) const override;
+  int DrawTextLines(int x, int y, const std::vector<std::string>& lines) const override;
+  int DrawWrappedTextLines(int x, int y, const std::vector<std::string>& lines) const override;
 
   Icon currentIcon;
 
diff --git a/tests/unit/screen_ui_test.cpp b/tests/unit/screen_ui_test.cpp
index 7d97a006b..ec269503e 100644
--- a/tests/unit/screen_ui_test.cpp
+++ b/tests/unit/screen_ui_test.cpp
@@ -38,8 +38,39 @@
 static const std::vector<std::string> HEADERS{ "header" };
 static const std::vector<std::string> ITEMS{ "item1", "item2", "item3", "item4", "1234567890" };
 
-TEST(ScreenUITest, StartPhoneMenuSmoke) {
-  Menu menu(false, 10, 20, HEADERS, ITEMS, 0);
+// TODO(xunchang) check if some draw functions are called when drawing menus.
+class MockDrawFunctions : public DrawInterface {
+  void SetColor(UIElement /* element */) const override {}
+  void DrawHighlightBar(int /* x */, int /* y */, int /* width */,
+                        int /* height */) const override {};
+  int DrawHorizontalRule(int /* y */) const override {
+    return 0;
+  };
+  int DrawTextLine(int /* x */, int /* y */, const std::string& /* line */,
+                   bool /* bold */) const override {
+    return 0;
+  };
+  void DrawSurface(GRSurface* /* surface */, int /* sx */, int /* sy */, int /* w */, int /* h */,
+                   int /* dx */, int /* dy */) const override {};
+  void DrawFill(int /* x */, int /* y */, int /* w */, int /* h */) const override {};
+  void DrawTextIcon(int /* x */, int /* y */, GRSurface* /* surface */) const override {};
+  int DrawTextLines(int /* x */, int /* y */,
+                    const std::vector<std::string>& /* lines */) const override {
+    return 0;
+  };
+  int DrawWrappedTextLines(int /* x */, int /* y */,
+                           const std::vector<std::string>& /* lines */) const override {
+    return 0;
+  };
+};
+
+class ScreenUITest : public testing::Test {
+ protected:
+  MockDrawFunctions draw_funcs_;
+};
+
+TEST_F(ScreenUITest, StartPhoneMenuSmoke) {
+  TextMenu menu(false, 10, 20, HEADERS, ITEMS, 0, 20, draw_funcs_);
   ASSERT_FALSE(menu.scrollable());
   ASSERT_EQ(HEADERS[0], menu.text_headers()[0]);
   ASSERT_EQ(5u, menu.ItemsCount());
@@ -53,8 +84,8 @@ TEST(ScreenUITest, StartPhoneMenuSmoke) {
   ASSERT_EQ(0, menu.selection());
 }
 
-TEST(ScreenUITest, StartWearMenuSmoke) {
-  Menu menu(true, 10, 8, HEADERS, ITEMS, 1);
+TEST_F(ScreenUITest, StartWearMenuSmoke) {
+  TextMenu menu(true, 10, 8, HEADERS, ITEMS, 1, 20, draw_funcs_);
   ASSERT_TRUE(menu.scrollable());
   ASSERT_EQ(HEADERS[0], menu.text_headers()[0]);
   ASSERT_EQ(5u, menu.ItemsCount());
@@ -69,8 +100,8 @@ TEST(ScreenUITest, StartWearMenuSmoke) {
   ASSERT_EQ(1, menu.selection());
 }
 
-TEST(ScreenUITest, StartPhoneMenuItemsOverflow) {
-  Menu menu(false, 1, 20, HEADERS, ITEMS, 0);
+TEST_F(ScreenUITest, StartPhoneMenuItemsOverflow) {
+  TextMenu menu(false, 1, 20, HEADERS, ITEMS, 0, 20, draw_funcs_);
   ASSERT_FALSE(menu.scrollable());
   ASSERT_EQ(1u, menu.ItemsCount());
 
@@ -84,8 +115,8 @@ TEST(ScreenUITest, StartPhoneMenuItemsOverflow) {
   ASSERT_EQ(1u, menu.MenuEnd());
 }
 
-TEST(ScreenUITest, StartWearMenuItemsOverflow) {
-  Menu menu(true, 1, 20, HEADERS, ITEMS, 0);
+TEST_F(ScreenUITest, StartWearMenuItemsOverflow) {
+  TextMenu menu(true, 1, 20, HEADERS, ITEMS, 0, 20, draw_funcs_);
   ASSERT_TRUE(menu.scrollable());
   ASSERT_EQ(5u, menu.ItemsCount());
 
@@ -101,9 +132,9 @@ TEST(ScreenUITest, StartWearMenuItemsOverflow) {
   ASSERT_EQ(1u, menu.MenuEnd());
 }
 
-TEST(ScreenUITest, PhoneMenuSelectSmoke) {
+TEST_F(ScreenUITest, PhoneMenuSelectSmoke) {
   int sel = 0;
-  Menu menu(false, 10, 20, HEADERS, ITEMS, sel);
+  TextMenu menu(false, 10, 20, HEADERS, ITEMS, sel, 20, draw_funcs_);
   // Mimic down button 10 times (2 * items size)
   for (int i = 0; i < 10; i++) {
     sel = menu.Select(++sel);
@@ -130,9 +161,9 @@ TEST(ScreenUITest, PhoneMenuSelectSmoke) {
   }
 }
 
-TEST(ScreenUITest, WearMenuSelectSmoke) {
+TEST_F(ScreenUITest, WearMenuSelectSmoke) {
   int sel = 0;
-  Menu menu(true, 10, 20, HEADERS, ITEMS, sel);
+  TextMenu menu(true, 10, 20, HEADERS, ITEMS, sel, 20, draw_funcs_);
   // Mimic pressing down button 10 times (2 * items size)
   for (int i = 0; i < 10; i++) {
     sel = menu.Select(++sel);
@@ -159,9 +190,9 @@ TEST(ScreenUITest, WearMenuSelectSmoke) {
   }
 }
 
-TEST(ScreenUITest, WearMenuSelectItemsOverflow) {
+TEST_F(ScreenUITest, WearMenuSelectItemsOverflow) {
   int sel = 1;
-  Menu menu(true, 3, 20, HEADERS, ITEMS, sel);
+  TextMenu menu(true, 3, 20, HEADERS, ITEMS, sel, 20, draw_funcs_);
   ASSERT_EQ(5u, menu.ItemsCount());
 
   // Scroll the menu to the end, and check the start & end of menu.
diff --git a/wear_ui.cpp b/wear_ui.cpp
index 3b057b761..8f3bc7bbe 100644
--- a/wear_ui.cpp
+++ b/wear_ui.cpp
@@ -73,7 +73,7 @@ void WearRecoveryUI::draw_screen_locked() {
   if (!show_text) {
     draw_foreground_locked();
   } else {
-    SetColor(TEXT_FILL);
+    SetColor(UIElement::TEXT_FILL);
     gr_fill(0, 0, gr_fb_width(), gr_fb_height());
 
     // clang-format off
@@ -99,8 +99,9 @@ void WearRecoveryUI::StartMenu(const std::vector<std::string>& headers,
                                const std::vector<std::string>& items, size_t initial_selection) {
   std::lock_guard<std::mutex> lg(updateMutex);
   if (text_rows_ > 0 && text_cols_ > 0) {
-    menu_ = std::make_unique<Menu>(scrollable_menu_, text_rows_ - menu_unusable_rows_ - 1,
-                                   text_cols_ - 1, headers, items, initial_selection);
+    menu_ = std::make_unique<TextMenu>(scrollable_menu_, text_rows_ - menu_unusable_rows_ - 1,
+                                       text_cols_ - 1, headers, items, initial_selection,
+                                       char_height_, *this);
     update_screen_locked();
   }
 }