From fe582c8b9848a7dcef39765344e32f3de946f7a8 Mon Sep 17 00:00:00 2001
From: John Wilkie <124276291+JBWilkie@users.noreply.github.com>
Date: Tue, 27 Aug 2024 16:43:03 +0100
Subject: [PATCH] [DAR-3333][External] Import raster layer annotations even if
 classes are created or updated (#918)

* Always return  when fetching remote classes

* Unit tests

* More concise tests
---
 darwin/dataset/remote_dataset.py            |  2 +
 darwin/importer/importer.py                 |  6 +-
 tests/darwin/dataset/remote_dataset_test.py | 73 +++++++++++++++++++++
 3 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/darwin/dataset/remote_dataset.py b/darwin/dataset/remote_dataset.py
index 154c26c85..e95cc00e9 100644
--- a/darwin/dataset/remote_dataset.py
+++ b/darwin/dataset/remote_dataset.py
@@ -691,6 +691,8 @@ def fetch_remote_classes(self, team_wide=False) -> List[Dict[str, Any]]:
             cls["available"] = belongs_to_current_dataset
             if team_wide or belongs_to_current_dataset:
                 classes_to_return.append(cls)
+            elif cls["annotation_types"] == ["raster_layer"]:
+                classes_to_return.append(cls)
         return classes_to_return
 
     def fetch_remote_attributes(self) -> List[Dict[str, Any]]:
diff --git a/darwin/importer/importer.py b/darwin/importer/importer.py
index d56a3eedd..138769520 100644
--- a/darwin/importer/importer.py
+++ b/darwin/importer/importer.py
@@ -1351,8 +1351,10 @@ def _import_annotations(
         )
 
         if (
-            annotation_type not in remote_classes
-            or annotation_class.name not in remote_classes[annotation_type]
+            (
+                annotation_type not in remote_classes
+                or annotation_class.name not in remote_classes[annotation_type]
+            )
             and annotation_type
             != "raster_layer"  # We do not skip raster layers as they are always available.
         ):
diff --git a/tests/darwin/dataset/remote_dataset_test.py b/tests/darwin/dataset/remote_dataset_test.py
index c974ace0e..7ed48a448 100644
--- a/tests/darwin/dataset/remote_dataset_test.py
+++ b/tests/darwin/dataset/remote_dataset_test.py
@@ -583,6 +583,79 @@ def test_fetches_files_with_commas(
         )
 
 
+@pytest.mark.usefixtures("file_read_write_test")
+class TestFetchRemoteClasses:
+    def setup_method(self):
+        self.mock_classes = [
+            {
+                "name": "class1",
+                "datasets": [{"id": 1}],
+                "annotation_types": ["type1"],
+            },
+            {
+                "name": "class2",
+                "datasets": [{"id": 2}],
+                "annotation_types": ["type2"],
+            },
+            {
+                "name": "raster_class",
+                "datasets": [],
+                "annotation_types": ["raster_layer"],
+            },
+        ]
+
+    def create_remote_dataset(
+        self, darwin_client, dataset_name, dataset_slug, team_slug_darwin_json_v2
+    ):
+        return RemoteDatasetV2(
+            client=darwin_client,
+            team=team_slug_darwin_json_v2,
+            name=dataset_name,
+            slug=dataset_slug,
+            dataset_id=1,
+        )
+
+    @responses.activate
+    def test_fetch_remote_classes_team_wide(
+        self,
+        darwin_client: Client,
+        dataset_name: str,
+        dataset_slug: str,
+        team_slug_darwin_json_v2: str,
+    ):
+        remote_dataset = self.create_remote_dataset(
+            darwin_client, dataset_name, dataset_slug, team_slug_darwin_json_v2
+        )
+        with patch.object(
+            remote_dataset.client,
+            "fetch_remote_classes",
+            return_value=self.mock_classes,
+        ):
+            result = remote_dataset.fetch_remote_classes(team_wide=True)
+            assert len(result) == 3
+            assert any(cls["name"] == "raster_class" for cls in result)
+
+    @responses.activate
+    def test_fetch_remote_classes_local_to_dataset(
+        self,
+        darwin_client: Client,
+        dataset_name: str,
+        dataset_slug: str,
+        team_slug_darwin_json_v2: str,
+    ):
+        remote_dataset = self.create_remote_dataset(
+            darwin_client, dataset_name, dataset_slug, team_slug_darwin_json_v2
+        )
+        with patch.object(
+            remote_dataset.client,
+            "fetch_remote_classes",
+            return_value=self.mock_classes,
+        ):
+            result = remote_dataset.fetch_remote_classes(team_wide=False)
+            assert len(result) == 2
+            assert any(cls["name"] == "raster_class" for cls in result)
+
+
 @pytest.fixture
 def remote_dataset(
     darwin_client: Client,