diff --git a/content/treestyletab/treestyletab.js b/content/treestyletab/treestyletab.js
index 5b919076..92312ff9 100644
--- a/content/treestyletab/treestyletab.js
+++ b/content/treestyletab/treestyletab.js
@@ -648,16 +648,19 @@ var TreeStyleTabService = {
 			).singleNodeValue;
 	},
  
-	getTabBrowserFromChild : function(aTab) 
+	getTabBrowserFromChild : function(aTabBrowserChild) 
 	{
-		if (!aTab) return null;
+		if (!aTabBrowserChild) return null;
 
-		if (aTab.__treestyletab__linkedTabBrowser)
-			return aTab.__treestyletab__linkedTabBrowser;
+		if (aTabBrowserChild.__treestyletab__linkedTabBrowser)
+			return aTabBrowserChild.__treestyletab__linkedTabBrowser;
+
+		if (aTabBrowserChild.localName == 'tabbrowser')
+			return aTabBrowserChild;
 
 		return this.evaluateXPath(
-				'ancestor-or-self::xul:tabbrowser',
-				aTab,
+				'ancestor::xul:tabbrowser[1]',
+				aTabBrowserChild,
 				XPathResult.FIRST_ORDERED_NODE_TYPE
 			).singleNodeValue;
 	},
@@ -751,17 +754,18 @@ var TreeStyleTabService = {
 			).booleanValue;
 	},
  
-	getTabs : function(aTabBrowser) 
+	getTabs : function(aTabBrowserChild) 
 	{
+		var b = this.getTabBrowserFromChild(aTabBrowserChild);
 		return this.evaluateXPath(
 				'descendant::xul:tab',
-				aTabBrowser.mTabContainer
+				b.mTabContainer
 			);
 	},
  
-	getTabsArray : function(aTabBrowser) 
+	getTabsArray : function(aTabBrowserChild) 
 	{
-		var tabs = this.getTabs(aTabBrowser);
+		var tabs = this.getTabs(aTabBrowserChild);
 		var array = [];
 		for (var i = 0, maxi = tabs.snapshotLength; i < maxi; i++)
 		{
@@ -770,20 +774,22 @@ var TreeStyleTabService = {
 		return array;
 	},
  
-	getFirstTab : function(aTabBrowser) 
+	getFirstTab : function(aTabBrowserChild) 
 	{
+		var b = this.getTabBrowserFromChild(aTabBrowserChild);
 		return this.evaluateXPath(
-				'descendant::xul:tab[1]',
-				aTabBrowser.mTabContainer,
+				'child::xul:tab[1]',
+				b.mTabContainer,
 				XPathResult.FIRST_ORDERED_NODE_TYPE
 			).singleNodeValue;
 	},
  
-	getLastTab : function(aTabBrowser) 
+	getLastTab : function(aTabBrowserChild) 
 	{
+		var b = this.getTabBrowserFromChild(aTabBrowserChild);
 		return this.evaluateXPath(
-				'descendant::xul:tab[last()]',
-				aTabBrowser.mTabContainer,
+				'child::xul:tab[last()]',
+				b.mTabContainer,
 				XPathResult.FIRST_ORDERED_NODE_TYPE
 			).singleNodeValue;
 	},
@@ -846,28 +852,30 @@ var TreeStyleTabService = {
 			).singleNodeValue;
 	},
  
-	getLastVisibleTab : function(aTab) 
+	getLastVisibleTab : function(aTabBrowserChild) 
 	{
-		if (!aTab) return null;
+		var b = this.getTabBrowserFromChild(aTabBrowserChild);
+		if (!b) return null;
 
-		if (!this.canCollapseSubtree(aTab))
-			return this.getLastTab(aTab);
+		if (!this.canCollapseSubtree(b))
+			return this.getLastTab(b);
 
 		return this.evaluateXPath(
 				'child::xul:tab[not(@'+this.kCOLLAPSED+'="true")][last()]',
-				aTab.parentNode,
+				b.mTabContainer,
 				XPathResult.FIRST_ORDERED_NODE_TYPE
 			).singleNodeValue;
 	},
  
-	getVisibleTabs : function(aTab) 
+	getVisibleTabs : function(aTabBrowserChild) 
 	{
-		if (!this.canCollapseSubtree(aTab))
+		var b = this.getTabBrowserFromChild(aTabBrowserChild);
+		if (!this.canCollapseSubtree(b))
 			return this.getTabs(b);
 
 		var xpathResult = this.evaluateXPath(
 				'child::xul:tab[not(@'+this.kCOLLAPSED+'="true")]',
-				aTab.parentNode
+				b.mTabContainer
 			);
 		return xpathResult;
 	},
diff --git a/tests/unit/treestyletab_tabUtils.test.js b/tests/unit/treestyletab_tabUtils.test.js
index 25286cb5..22bf7dc7 100644
--- a/tests/unit/treestyletab_tabUtils.test.js
+++ b/tests/unit/treestyletab_tabUtils.test.js
@@ -139,43 +139,136 @@ function test_getTabs()
 	assert.isNull(sv.getNextTab(tabs[3]));
 	assert.equals(tabs[1], sv.getPreviousTab(tabs[2]));
 	assert.isNull(sv.getPreviousTab(tabs[0]));
+
+	assert.equals(0, sv.getTabIndex(tabs[0]));
+	assert.equals(1, sv.getTabIndex(tabs[1]));
+	assert.equals(2, sv.getTabIndex(tabs[2]));
+	assert.equals(3, sv.getTabIndex(tabs[3]));
 }
 
-function test_tabsVisibility()
+function test_canCollapseSubtree()
+{
+	gBrowser.setAttribute(sv.kALLOW_COLLAPSE, true);
+	assert.isTrue(sv.canCollapseSubtree(gBrowser));
+
+	gBrowser.removeAttribute(sv.kALLOW_COLLAPSE);
+	assert.isFalse(sv.canCollapseSubtree(gBrowser));
+}
+
+function test_isCollapsed()
 {
 	tabs[1].setAttribute(sv.kCOLLAPSED, true);
 	tabs[3].setAttribute(sv.kCOLLAPSED, true);
 
+	gBrowser.setAttribute(sv.kALLOW_COLLAPSE, true);
+	assert.isFalse(sv.isCollapsed(tabs[0]));
+	assert.isTrue(sv.isCollapsed(tabs[1]));
+	assert.isFalse(sv.isCollapsed(tabs[2]));
+	assert.isTrue(sv.isCollapsed(tabs[3]));
+
+	gBrowser.removeAttribute(sv.kALLOW_COLLAPSE);
+	assert.isFalse(sv.isCollapsed(tabs[0]));
+	assert.isFalse(sv.isCollapsed(tabs[1]));
+	assert.isFalse(sv.isCollapsed(tabs[2]));
+	assert.isFalse(sv.isCollapsed(tabs[3]));
+}
+
+function test_getNextVisibleTab()
+{
+	tabs[1].setAttribute(sv.kCOLLAPSED, true);
+	tabs[3].setAttribute(sv.kCOLLAPSED, true);
+
+	gBrowser.setAttribute(sv.kALLOW_COLLAPSE, true);
 	assert.equals(tabs[2], sv.getNextVisibleTab(tabs[0]));
 	assert.equals(tabs[2], sv.getNextVisibleTab(tabs[1]));
 	assert.isNull(sv.getNextVisibleTab(tabs[2]));
 	assert.isNull(sv.getNextVisibleTab(tabs[3]));
+
+	gBrowser.removeAttribute(sv.kALLOW_COLLAPSE);
+	assert.equals(tabs[1], sv.getNextVisibleTab(tabs[0]));
+	assert.equals(tabs[2], sv.getNextVisibleTab(tabs[1]));
+	assert.equals(tabs[3], sv.getNextVisibleTab(tabs[2]));
+	assert.isNull(sv.getNextVisibleTab(tabs[3]));
+}
+
+function test_getPreviousVisibleTab()
+{
+	tabs[1].setAttribute(sv.kCOLLAPSED, true);
+	tabs[3].setAttribute(sv.kCOLLAPSED, true);
+
+	gBrowser.setAttribute(sv.kALLOW_COLLAPSE, true);
 	assert.isNull(sv.getPreviousVisibleTab(tabs[0]));
 	assert.equals(tabs[0], sv.getPreviousVisibleTab(tabs[1]));
 	assert.equals(tabs[0], sv.getPreviousVisibleTab(tabs[2]));
 	assert.equals(tabs[2], sv.getPreviousVisibleTab(tabs[3]));
 
+	gBrowser.removeAttribute(sv.kALLOW_COLLAPSE);
+	assert.equals(tabs[1], sv.getNextVisibleTab(tabs[0]));
+	assert.equals(tabs[2], sv.getNextVisibleTab(tabs[1]));
+	assert.equals(tabs[3], sv.getNextVisibleTab(tabs[2]));
+	assert.isNull(sv.getNextVisibleTab(tabs[3]));
+}
+
+function test_getLastVisibleTab()
+{
+	tabs[1].setAttribute(sv.kCOLLAPSED, true);
+	tabs[3].setAttribute(sv.kCOLLAPSED, true);
+
+	gBrowser.setAttribute(sv.kALLOW_COLLAPSE, true);
 	assert.equals(tabs[2], sv.getLastVisibleTab(tabs[0]));
 	assert.equals(tabs[2], sv.getLastVisibleTab(tabs[1]));
 	assert.equals(tabs[2], sv.getLastVisibleTab(tabs[2]));
 	assert.equals(tabs[2], sv.getLastVisibleTab(tabs[3]));
 
+	gBrowser.removeAttribute(sv.kALLOW_COLLAPSE);
+	assert.equals(tabs[3], sv.getLastVisibleTab(tabs[0]));
+	assert.equals(tabs[3], sv.getLastVisibleTab(tabs[1]));
+	assert.equals(tabs[3], sv.getLastVisibleTab(tabs[2]));
+	assert.equals(tabs[3], sv.getLastVisibleTab(tabs[3]));
+}
+
+function test_getVisibleTabs()
+{
+	tabs[1].setAttribute(sv.kCOLLAPSED, true);
+	tabs[3].setAttribute(sv.kCOLLAPSED, true);
+
+	gBrowser.setAttribute(sv.kALLOW_COLLAPSE, true);
+
 	var visibleResult = sv.getVisibleTabs(tabs[0]);
-	assert.isTrue(visibleResult instanceof XPathResult);
+	assert.implementsInterface(Ci.nsIDOMXPathResult, visibleResult);
 	assert.equals(2, visibleResult.snapshotLength);
 
 	var visibleTabs = [];
-	for (var i = 0, maxi = visibleResult.snapshotLength; i < maxi; i++)
+	for (let i = 0, maxi = visibleResult.snapshotLength; i < maxi; i++)
 	{
 		visibleTabs.push(visibleResult.snapshotItem(i));
 	}
 	assert.equals(2, visibleTabs.length);
 	assert.equals([tabs[0], tabs[2]], visibleTabs);
 
+	gBrowser.removeAttribute(sv.kALLOW_COLLAPSE);
+
+	visibleResult = sv.getVisibleTabs(tabs[0]);
+	assert.implementsInterface(Ci.nsIDOMXPathResult, visibleResult);
+	assert.equals(4, visibleResult.snapshotLength);
+}
+
+function test_getVisibleIndex()
+{
+	tabs[1].setAttribute(sv.kCOLLAPSED, true);
+	tabs[3].setAttribute(sv.kCOLLAPSED, true);
+
+	gBrowser.setAttribute(sv.kALLOW_COLLAPSE, true);
 	assert.equals(0, sv.getVisibleIndex(tabs[0]));
 	assert.equals(-1, sv.getVisibleIndex(tabs[1]));
 	assert.equals(1, sv.getVisibleIndex(tabs[2]));
 	assert.equals(-1, sv.getVisibleIndex(tabs[3]));
+
+	gBrowser.removeAttribute(sv.kALLOW_COLLAPSE);
+	assert.equals(0, sv.getVisibleIndex(tabs[0]));
+	assert.equals(1, sv.getVisibleIndex(tabs[1]));
+	assert.equals(2, sv.getVisibleIndex(tabs[2]));
+	assert.equals(3, sv.getVisibleIndex(tabs[3]));
 }
 
 var randomKey = 'key-'+parseInt(Math.random() * 65000);