From e2888f12d7703f432ec69a5a20734eb698bb6908 Mon Sep 17 00:00:00 2001
From: shamoon <4887959+shamoon@users.noreply.github.com>
Date: Sat, 1 Mar 2025 19:54:12 -0800
Subject: [PATCH] Fix some handling, add tests

---
 src/documents/signals/handlers.py     | 41 ++++++++++++++++-----------
 src/documents/tests/test_workflows.py | 16 +++++++++--
 2 files changed, 38 insertions(+), 19 deletions(-)

diff --git a/src/documents/signals/handlers.py b/src/documents/signals/handlers.py
index 342167446..4eb7f72cc 100644
--- a/src/documents/signals/handlers.py
+++ b/src/documents/signals/handlers.py
@@ -576,8 +576,6 @@ def cleanup_custom_field_deletion(sender, instance: CustomField, **kwargs):
             f"Removing custom field {instance} from sort field of {views_with_sort_updated} views",
         )
 
-    # Remove from workflow actions
-
 
 def add_to_index(sender, document, **kwargs):
     from documents import index
@@ -772,29 +770,38 @@ def run_workflows(
         if action.assign_custom_fields.exists():
             if not use_overrides:
                 for field in action.assign_custom_fields.all():
-                    if not CustomFieldInstance.objects.filter(
+                    value_field_name = CustomFieldInstance.get_value_field_name(
+                        data_type=field.data_type,
+                    )
+                    args = {
+                        value_field_name: action.assign_custom_fields_values.get(
+                            str(field.pk),
+                            None,
+                        ),
+                    }
+                    # for some reason update_or_create doesn't work here
+                    instance = CustomFieldInstance.objects.filter(
                         field=field,
                         document=document,
-                    ).exists():
-                        # can be triggered on existing docs, so only add the field if it doesn't already exist
-                        value_field_name = CustomFieldInstance.get_value_field_name(
-                            data_type=field.data_type,
+                    ).first()
+                    if instance:
+                        setattr(instance, value_field_name, args[value_field_name])
+                        instance.save()
+                    else:
+                        CustomFieldInstance.objects.create(
+                            **args,
+                            field=field,
+                            document=document,
                         )
-                        args = {
-                            "field": field,
-                            "document": document,
-                            value_field_name: action.assign_custom_fields_values.get(
-                                field.pk,
-                                None,
-                            ),
-                        }
-                        CustomFieldInstance.objects.create(**args)
             else:
                 if overrides.custom_fields is None:
                     overrides.custom_fields = {}
                 overrides.custom_fields.update(
                     {
-                        field.pk: action.assign_custom_fields_values.get(field.pk, None)
+                        field.pk: action.assign_custom_fields_values.get(
+                            str(field.pk),
+                            None,
+                        )
                         for field in action.assign_custom_fields.all()
                     },
                 )
diff --git a/src/documents/tests/test_workflows.py b/src/documents/tests/test_workflows.py
index 94dcb7689..3006594cc 100644
--- a/src/documents/tests/test_workflows.py
+++ b/src/documents/tests/test_workflows.py
@@ -133,6 +133,9 @@ class TestWorkflows(
         action.assign_change_groups.add(self.group1.pk)
         action.assign_custom_fields.add(self.cf1.pk)
         action.assign_custom_fields.add(self.cf2.pk)
+        action.assign_custom_fields_values = {
+            self.cf2.pk: 42,
+        }
         action.save()
         w = Workflow.objects.create(
             name="Workflow 1",
@@ -209,6 +212,10 @@ class TestWorkflows(
                     list(document.custom_fields.all().values_list("field", flat=True)),
                     [self.cf1.pk, self.cf2.pk],
                 )
+                self.assertEqual(
+                    document.custom_fields.get(field=self.cf2.pk).value,
+                    42,
+                )
 
         info = cm.output[0]
         expected_str = f"Document matched {trigger} from {w}"
@@ -1215,11 +1222,11 @@ class TestWorkflows(
     def test_document_updated_workflow_existing_custom_field(self):
         """
         GIVEN:
-            - Existing workflow with UPDATED trigger and action that adds a custom field
+            - Existing workflow with UPDATED trigger and action that assigns a custom field with a value
         WHEN:
             - Document is updated that already contains the field
         THEN:
-            - Document update succeeds without trying to re-create the field
+            - Document update succeeds and updates the field
         """
         trigger = WorkflowTrigger.objects.create(
             type=WorkflowTrigger.WorkflowTriggerType.DOCUMENT_UPDATED,
@@ -1227,6 +1234,8 @@ class TestWorkflows(
         )
         action = WorkflowAction.objects.create()
         action.assign_custom_fields.add(self.cf1)
+        action.assign_custom_fields_values = {self.cf1.pk: "new value"}
+        action.save()
         w = Workflow.objects.create(
             name="Workflow 1",
             order=0,
@@ -1251,6 +1260,9 @@ class TestWorkflows(
             format="json",
         )
 
+        doc.refresh_from_db()
+        self.assertEqual(doc.custom_fields.get(field=self.cf1).value, "new value")
+
     def test_document_updated_workflow_merge_permissions(self):
         """
         GIVEN: