diff --git a/tasks/assignment/assets/__init__.py b/tasks/assignment/assets/__init__.py
deleted file mode 100644
index 92bc0ebf3..000000000
--- a/tasks/assignment/assets/__init__.py
+++ /dev/null
@@ -1,13 +0,0 @@
-from tasks.assignment.assets.assets_assignment_ui import (
-    ALL_ABOUT_BOOTHILL_CHECK, ALL_ABOUT_BOOTHILL_CLICK,
-    CHARACTER_MATERIALS_CHECK, CHARACTER_MATERIALS_CLICK,
-    EXP_MATERIALS_CREDITS_CHECK, EXP_MATERIALS_CREDITS_CLICK, GROUP_SEARCH,
-    SYNTHESIS_MATERIALS_CHECK, SYNTHESIS_MATERIALS_CLICK)
-
-for group_button_wrapper in (
-    ALL_ABOUT_BOOTHILL_CHECK, ALL_ABOUT_BOOTHILL_CLICK,
-    CHARACTER_MATERIALS_CHECK, CHARACTER_MATERIALS_CLICK,
-    EXP_MATERIALS_CREDITS_CHECK, EXP_MATERIALS_CREDITS_CLICK,
-    SYNTHESIS_MATERIALS_CHECK, SYNTHESIS_MATERIALS_CLICK,
-):
-    group_button_wrapper.load_search(GROUP_SEARCH.button)
diff --git a/tasks/assignment/ui.py b/tasks/assignment/ui.py
index 5eeba0732..f6eb29fe6 100644
--- a/tasks/assignment/ui.py
+++ b/tasks/assignment/ui.py
@@ -15,7 +15,7 @@ from tasks.assignment.assets.assets_assignment_ui import *
 from tasks.assignment.keywords import *
 from tasks.base.assets.assets_base_page import ASSIGNMENT_CHECK
 from tasks.base.ui import UI
-from tasks.dungeon.ui import DungeonTabSwitch as Switch
+from tasks.dungeon.ui import DungeonTabSwitch
 
 
 class AssignmentStatus(Enum):
@@ -102,7 +102,11 @@ class AssignmentOcr(Ocr):
         return matched
 
 
-ASSIGNMENT_GROUP_SWITCH = Switch(
+class AssignmentGroupSwitch(DungeonTabSwitch):
+    SEARCH_BUTTON = GROUP_SEARCH
+
+
+ASSIGNMENT_GROUP_SWITCH = AssignmentGroupSwitch(
     'AssignmentGroupSwitch',
     is_selector=True
 )
diff --git a/tasks/dungeon/ui.py b/tasks/dungeon/ui.py
index 35b45e17d..5faecc52d 100644
--- a/tasks/dungeon/ui.py
+++ b/tasks/dungeon/ui.py
@@ -33,12 +33,14 @@ from tasks.map.keywords import KEYWORDS_MAP_WORLD, MapPlane
 
 
 class DungeonTabSwitch(Switch):
+    SEARCH_BUTTON = TAB_SEARCH
+
     def add_state(self, state, check_button, click_button=None):
         # Load search
         if check_button is not None:
-            check_button.load_search(TAB_SEARCH.area)
+            check_button.load_search(self.__class__.SEARCH_BUTTON.area)
         if click_button is not None:
-            click_button.load_search(TAB_SEARCH.area)
+            click_button.load_search(self.__class__.SEARCH_BUTTON.area)
         return super().add_state(state, check_button, click_button)
 
     def click(self, state, main):