From e5a4dd9c00f02a7d4237867db77862a4eef807e0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gregor=20Jer=C5=A1e?= <gregor.jerse@genialis.com>
Date: Wed, 8 Nov 2023 09:26:21 +0100
Subject: [PATCH] Fix move_to_entity on data object

When entity is in collection also assign collection to the data object.
---
 docs/CHANGELOG.rst                |  3 +++
 resolwe/flow/models/data.py       |  7 +++++--
 resolwe/flow/tests/test_models.py | 14 ++++++++++++++
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/docs/CHANGELOG.rst b/docs/CHANGELOG.rst
index b45ba7e89..14eed20a3 100644
--- a/docs/CHANGELOG.rst
+++ b/docs/CHANGELOG.rst
@@ -38,6 +38,9 @@ Fixed
 - Create history object for ``Collection`` if none exists
 - Use the same postfix for all kubernetes volumes for the given data object
 - Fix ``handle_progress`` handler in the listener
+- Fix moving data to entity in collection also sets collection on the data
+  object
+
 
 
 ===================
diff --git a/resolwe/flow/models/data.py b/resolwe/flow/models/data.py
index 8ec6b7e36..77351a3ae 100644
--- a/resolwe/flow/models/data.py
+++ b/resolwe/flow/models/data.py
@@ -630,8 +630,11 @@ def move_to_entity(self, entity):
             raise ValidationError("Data object must not be removed from container.")
         self.entity = entity
         if entity:
-            self.permission_group = entity.permission_group
-        self.tags = entity.tags
+            if entity.collection:
+                self.move_to_collection(entity.collection)
+            else:
+                self.permission_group = entity.permission_group
+                self.tags = entity.tags
         self.save(update_fields=["permission_group", "entity", "tags"])
 
     def validate_change_collection(self, collection):
diff --git a/resolwe/flow/tests/test_models.py b/resolwe/flow/tests/test_models.py
index 7dfece73d..01795739c 100644
--- a/resolwe/flow/tests/test_models.py
+++ b/resolwe/flow/tests/test_models.py
@@ -138,6 +138,20 @@ def test_delete_chunked(self):
         Data.objects.all().delete_chunked(chunk_size=10)
         self.assertFalse(Data.objects.exists())
 
+    def test_move_to_entity(self):
+        # Create data outside container and move it to the entity in the collection.
+        process = Process.objects.create(contributor=self.contributor)
+        data = Data.objects.create(contributor=self.contributor, process=process)
+        self.assertIsNone(data.entity)
+        self.assertIsNone(data.collection)
+        collection = Collection.objects.create(contributor=self.contributor)
+        entity = Entity.objects.create(
+            contributor=self.contributor, collection=collection
+        )
+        data.move_to_entity(entity)
+        self.assertEqual(data.entity, entity)
+        self.assertEqual(data.collection, collection)
+
     def test_trim_name(self):
         process = Process.objects.create(
             contributor=self.contributor, data_name="test" * 50