Refactor

Author: dandavison

tests/test_client.py

@@ -1345,24 +1345,22 @@ def update_search_attributes(
             print("empty-but-non-none-clears")
             # Pass empty but non-None to clear all attributes
             return ScheduleUpdate(
-                input.description.schedule, search_attributes=TypedSearchAttributes([])
+                input.description.schedule,
+                search_attributes=TypedSearchAttributes.empty,
             )
         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"),
-                    ]
+                search_attributes=input.description.typed_search_attributes.updated(
+                    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(
@@ -1376,7 +1374,7 @@ def update_search_attributes(
 
     await handle.update(update_search_attributes)
 
-    async def none_is_noop_expectations():
+    async def none_is_noop_expectations() -> bool:
         desc = await handle.describe()
         return (
             desc.search_attributes[key_1.name] == [val_1]
@@ -1385,16 +1383,18 @@ async def none_is_noop_expectations():
             and desc.typed_search_attributes[key_2] == val_2
         )
 
-    async def empty_but_non_none_clears_expectations():
+    async def empty_but_non_none_clears_expectations() -> bool:
         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
+        return_ = (
+            desc.search_attributes[key_1.name] == [],
+            desc.search_attributes[key_2.name] == [],
+            len(desc.typed_search_attributes) == 0,
+            len(desc.search_attributes) == 0,
         )
+        print(return_)
+        return all(return_)
 
-    async def all_new_values_overwrite_expectations():
+    async def all_new_values_overwrite_expectations() -> bool:
         desc = await handle.describe()
         return (
             desc.search_attributes[key_1.name] == [val_1 + "-new"]
@@ -1403,7 +1403,7 @@ async def all_new_values_overwrite_expectations():
             and desc.typed_search_attributes[key_2] == val_2 + "-new"
         )
 
-    async def partial_new_values_overwrites_and_drops_expectations():
+    async def partial_new_values_overwrites_and_drops_expectations() -> bool:
         desc = await handle.describe()
         return (
             desc.search_attributes[key_1.name] == [val_1 + "-new"]