From 3982170886b5adfa5295efdf64dd5472e6ff1f1c Mon Sep 17 00:00:00 2001
From: Frans de Jonge <fransdejonge@gmail.com>
Date: Mon, 3 Apr 2017 23:45:51 +0200
Subject: [PATCH] MenuSorter: fix findById

---
 frontend/ui/menusorter.lua    | 22 ++++++++++++++--------
 spec/unit/menusorter_spec.lua |  7 ++++++-
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/frontend/ui/menusorter.lua b/frontend/ui/menusorter.lua
index 5875d2a4a..85074f4bb 100644
--- a/frontend/ui/menusorter.lua
+++ b/frontend/ui/menusorter.lua
@@ -103,9 +103,11 @@ function MenuSorter:sort(item_table, order)
     for i,sub_menu in ipairs(sub_menus) do
         local sub_menu_position = self:findById(menu_table["KOMenu:menu_buttons"], sub_menu)
         if sub_menu_position then
-            sub_menu_position.sub_item_table = menu_table[sub_menu]
+            local sub_menu_content = menu_table[sub_menu]
+            sub_menu_position.text = sub_menu_content.text
+            sub_menu_position.sub_item_table = sub_menu_content
             -- remove reference from top level output
-            menu_table[sub_menu] = nil
+            sub_menu_content = nil
             -- remove reference from input so it won't show up as orphaned
             item_table[sub_menu] = nil
         end
@@ -150,17 +152,21 @@ function MenuSorter:findById(tbl, needle_id)
     local items = {}
 
     for _,item in pairs(tbl) do
-        table.insert(items, item)
+        if item ~= "KOMenu:menu_buttons" then
+            table.insert(items, item)
+        end
     end
 
     local k, v
     k, v = next(items, nil)
     while k do
-        if type(k) == "number" or k == "sub_item_table" then
-            if v.id == needle_id then
-                return v
-            elseif type(v) == "table" and v.id then
-                table.insert(items, v)
+        if v.id == needle_id then
+            return v
+        elseif v.sub_item_table then
+            for _,item in pairs(v.sub_item_table) do
+                if type(item) == "table" and item.id then
+                    table.insert(items, item)
+                end
             end
         end
         k, v = next(items, k)
diff --git a/spec/unit/menusorter_spec.lua b/spec/unit/menusorter_spec.lua
index 13f1058de..05381c618 100644
--- a/spec/unit/menusorter_spec.lua
+++ b/spec/unit/menusorter_spec.lua
@@ -38,16 +38,21 @@ describe("MenuSorter module", function()
             ["KOMenu:menu_buttons"] = {},
             first = {},
             second = {},
+            third1 = {},
+            third2 = {},
         }
         local order = {
             ["KOMenu:menu_buttons"] = {"first",},
-            first = {"second",},
+            first = {"second"},
+            second = {"third1", "third2"},
         }
 
         local test_menu = MenuSorter:sort(menu_items, order)
 
         assert.is_true(test_menu[1].id == "first")
         assert.is_true(test_menu[1][1].id == "second")
+        assert.is_true(test_menu[1][1].sub_item_table[1].id == "third1")
+        assert.is_true(test_menu[1][1].sub_item_table[2].id == "third2")
     end)
     it("should put orphans in the first menu", function()
         local menu_items = {