WIP: tests

Author: dandavison

tests/test_client.py

@@ -1287,8 +1287,17 @@ def update_schedule_typed_attrs(
     assert desc.typed_search_attributes[text_attr_key] == "some-schedule-attr1"
 
 
+@pytest.mark.parametrize(
+    "test_case",
+    [
+        "none-is-noop",
+        "empty-but-non-none-clears",
+        "all-new-values-overwrites",
+        "partial-new-values-overwrites-and-drops",
+    ],
+)
 async def test_schedule_search_attribute_update(
-    client: Client, env: WorkflowEnvironment
+    client: Client, env: WorkflowEnvironment, test_case: str
 ):
     if env.supports_time_skipping:
         pytest.skip("Java test server doesn't support schedules")
@@ -1319,31 +1328,102 @@ async def test_schedule_search_attribute_update(
         ),
     )
 
-    # Do update of search attributes
-    def update_schedule_search_attributes(
+    def update_search_attributes(
         input: ScheduleUpdateInput,
     ) -> Optional[ScheduleUpdate]:
-        # Make sure the search attributes are present in all forms
+        # Make sure the initial search attributes are present
         assert input.description.search_attributes[key_1.name] == [val_1]
+        assert input.description.search_attributes[key_2.name] == [val_2]
         assert input.description.typed_search_attributes[key_1] == val_1
+        assert input.description.typed_search_attributes[key_2] == val_2
+
+        if test_case == "none-is-noop":
+            print("none-is-noop")
+            # Passing None makes no changes
+            return ScheduleUpdate(input.description.schedule, search_attributes=None)
+        elif test_case == "empty-but-non-none-clears":
+            print("empty-but-non-none-clears")
+            # Pass empty but non-None to clear all attributes
+            return ScheduleUpdate(
+                input.description.schedule, search_attributes=TypedSearchAttributes([])
+            )
+        elif test_case == "all-new-values-overwrites":
+            print("all-new-values-overwrites")
+            # Pass all new values to overwrite existing
+            return ScheduleUpdate(
+                input.description.schedule,
+                search_attributes=TypedSearchAttributes(
+                    [
+                        SearchAttributePair(key_1, val_1 + "-new"),
+                        SearchAttributePair(key_2, val_2 + "-new"),
+                    ]
+                ),
+            )
+        elif test_case == "partial-new-values-overwrites-and-drops":
+            print("partial-new-values-overwrites-and-drops")
+            # Only update key_1, which should drop key_2
+            input.description.typed_search_attributes.updated
+            return ScheduleUpdate(
+                input.description.schedule,
+                search_attributes=TypedSearchAttributes(
+                    [
+                        SearchAttributePair(key_1, val_1 + "-new"),
+                    ]
+                ),
+            )
+        else:
+            raise ValueError(f"Invalid test case: {test_case}")
 
-        return ScheduleUpdate(
-            input.description.schedule,
-            search_attributes=input.description.typed_search_attributes.updated(
-                SearchAttributePair(key_1, val_1 + "-updated")
-            ),
+    await handle.update(update_search_attributes)
+
+    async def none_is_noop_expectations():
+        desc = await handle.describe()
+        return (
+            desc.search_attributes[key_1.name] == [val_1]
+            and desc.search_attributes[key_2.name] == [val_2]
+            and desc.typed_search_attributes[key_1] == val_1
+            and desc.typed_search_attributes[key_2] == val_2
         )
 
-    await handle.update(update_schedule_search_attributes)
+    async def empty_but_non_none_clears_expectations():
+        desc = await handle.describe()
+        return (
+            desc.search_attributes[key_1.name] == []
+            and desc.search_attributes[key_2.name] == []
+            and len(desc.typed_search_attributes) == 0
+            and len(desc.search_attributes) == 0
+        )
 
-    async def check_fn():
+    async def all_new_values_overwrite_expectations():
         desc = await handle.describe()
         return (
-            desc.typed_search_attributes[key_1],
-            desc.search_attributes[key_1.name],
+            desc.search_attributes[key_1.name] == [val_1 + "-new"]
+            and desc.search_attributes[key_2.name] == [val_2 + "-new"]
+            and desc.typed_search_attributes[key_1] == val_1 + "-new"
+            and desc.typed_search_attributes[key_2] == val_2 + "-new"
         )
 
-    await assert_eq_eventually((val_1 + "-updated", [val_1 + "-updated"]), check_fn)
+    async def partial_new_values_overwrites_and_drops_expectations():
+        desc = await handle.describe()
+        return (
+            desc.search_attributes[key_1.name] == [val_1 + "-new"]
+            and desc.typed_search_attributes[key_1] == val_1 + "-new"
+            and key_2.name not in desc.search_attributes
+            and key_2 not in desc.typed_search_attributes
+        )
+
+    if test_case == "none-is-noop":  # 🟢
+        await assert_eq_eventually(True, none_is_noop_expectations)
+    elif test_case == "empty-but-non-none-clears":  # 🔴
+        await assert_eq_eventually(True, empty_but_non_none_clears_expectations)
+    elif test_case == "all-new-values-overwrites":  # 🟢
+        await assert_eq_eventually(True, all_new_values_overwrite_expectations)
+    elif test_case == "partial-new-values-overwrites-and-drops":  # 🟢
+        await assert_eq_eventually(
+            True, partial_new_values_overwrites_and_drops_expectations
+        )
+    else:
+        raise ValueError(f"Invalid test case: {test_case}")
 
 
 async def assert_no_schedules(client: Client) -> None: