From 4a513e1e41373c3bee6c1c00445370115a2fceed Mon Sep 17 00:00:00 2001
From: Andrej730 <azhilenkov@gmail.com>
Date: Mon, 25 Mar 2024 12:17:38 +0500
Subject: [PATCH] fix linked ifc visibility toggle tests

---
 .../blenderbim/bim/module/project/operator.py | 12 ++------
 src/blenderbim/blenderbim/tool/blender.py     | 29 +++++++++++++++++++
 .../test/bim/feature/project.feature          | 10 +++++--
 src/blenderbim/test/bim/test_feature.py       | 15 ++++++++++
 4 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/src/blenderbim/blenderbim/bim/module/project/operator.py b/src/blenderbim/blenderbim/bim/module/project/operator.py
index 54761de6f9..8dc58c4651 100644
--- a/src/blenderbim/blenderbim/bim/module/project/operator.py
+++ b/src/blenderbim/blenderbim/bim/module/project/operator.py
@@ -1063,17 +1063,9 @@ class ToggleLinkVisibility(bpy.types.Operator):
 
     def toggle_visibility(self, link):
         linked_collections = self.get_linked_collections()
-        queue = [bpy.context.view_layer.layer_collection]
-        layer_collection = None
 
-        while queue:
-            layer = queue.pop()
-            if layer.collection in linked_collections:
-                layer_collection = layer
-                break
-            queue.extend(list(layer.children))
-
-        if layer_collection:
+        layer_collections = tool.Blender.get_layer_collections_mapping(linked_collections)
+        for layer_collection in layer_collections.values():
             layer_collection.exclude = not layer_collection.exclude
             link.is_hidden = layer_collection.exclude
 
diff --git a/src/blenderbim/blenderbim/tool/blender.py b/src/blenderbim/blenderbim/tool/blender.py
index 73a7c7626a..a6de0de09a 100644
--- a/src/blenderbim/blenderbim/tool/blender.py
+++ b/src/blenderbim/blenderbim/tool/blender.py
@@ -648,6 +648,35 @@ class Blender(blenderbim.core.tool.Blender):
         open_file_or_folder(filepath.as_posix())
         return {"PASS_THROUGH"}
 
+    @classmethod
+    def get_layer_collection(
+        cls, collection: bpy.types.Collection, view_layer: Optional[bpy.types.ViewLayer] = None
+    ) -> Union[bpy.types.LayerCollection, None]:
+        return cls.get_layer_collections_mapping([collection], view_layer).get(collection)
+
+    @classmethod
+    def get_layer_collections_mapping(
+        cls, collections: list[bpy.types.Collection], view_layer: Optional[bpy.types.ViewLayer] = None
+    ) -> dict[bpy.types.Collection, bpy.types.LayerCollection]:
+        if view_layer is None:
+            view_layer = bpy.context.view_layer
+
+        collections = list(collections)  # copy to prevent mutation
+        collections_mapping = dict()
+        queue = [view_layer.layer_collection]
+
+        while queue:
+            layer = queue.pop()
+            collection = layer.collection
+            if collection in collections:
+                collections_mapping[collection] = layer
+                collections.remove(collection)
+                if not collections:
+                    break
+            queue.extend(list(layer.children))
+
+        return collections_mapping
+
     class Modifier:
         @classmethod
         def is_eligible_for_railing_modifier(cls, obj):
diff --git a/src/blenderbim/test/bim/feature/project.feature b/src/blenderbim/test/bim/feature/project.feature
index 256a9397db..bc1e738720 100644
--- a/src/blenderbim/test/bim/feature/project.feature
+++ b/src/blenderbim/test/bim/feature/project.feature
@@ -331,9 +331,13 @@ Scenario: Toggle link visibility - wireframe mode
 
 Scenario: Toggle link visibility - visible mode
     Given an empty IFC project
-    And I press "bim.link_ifc(filepath='{cwd}/test/files/basic.blend')"
-    When I press "bim.toggle_link_visibility(link='{cwd}/test/files/basic.blend', mode='VISIBLE')"
-    Then nothing happens
+    And I press "bim.link_ifc(filepath='{cwd}/test/files/basic.ifc')"
+    When I press "bim.toggle_link_visibility(link='{cwd}/test/files/basic.ifc', mode='VISIBLE')"
+    Then "scene.BIMProjectProperties.links['{cwd}/test/files/basic.ifc'].is_hidden" is "True"
+    And the collection "IfcProject/basic.ifc" exclude status is "True"
+    When I press "bim.toggle_link_visibility(link='{cwd}/test/files/basic.ifc', mode='VISIBLE')"
+    Then "scene.BIMProjectProperties.links['{cwd}/test/files/basic.ifc'].is_hidden" is "False"
+    And the collection "IfcProject/basic.ifc" exclude status is "False"
 
 Scenario: Unload link
     Given an empty Blender session
diff --git a/src/blenderbim/test/bim/test_feature.py b/src/blenderbim/test/bim/test_feature.py
index f2be9ccb89..e8ab73555c 100644
--- a/src/blenderbim/test/bim/test_feature.py
+++ b/src/blenderbim/test/bim/test_feature.py
@@ -397,6 +397,21 @@ def the_collection_name_exists(name) -> bpy.types.Collection:
     return obj
 
 
+@then(parsers.parse('the collection "{name}" exists in viewlayer'))
+def the_collection_exists_in_viewlayer(name: str) -> bpy.types.LayerCollection:
+    col = the_collection_name_exists(name)
+    layer = tool.Blender.get_layer_collection(col)
+    if not layer:
+        assert False, f'The collection "{name}" is not present in the current viewlayer'
+    return layer
+
+
+@then(parsers.parse('the collection "{name}" exclude status is "{exclude}"'))
+def the_collection_exclude_status_is(name: str, exclude: str) -> None:
+    layer = the_collection_exists_in_viewlayer(name)
+    assert layer.exclude == (exclude == "True")
+
+
 @then(parsers.parse('the object "{name1}" and "{name2}" are different elements'))
 def the_object_name1_and_name2_are_different_elements(name1, name2):
     ifc = an_ifc_file_exists()