Address PR review comments

Author: andystaples

durabletask/client.py

@@ -247,7 +247,7 @@ def list_instance_ids(self,
                           page_size: Optional[int] = None,
                           continuation_token: Optional[str] = None) -> Page[str]:
         req = pb.ListInstanceIdsRequest(
-            runtimeStatus=[status.value for status in runtime_status] if runtime_status else None,
+            runtimeStatus=[status.value for status in runtime_status] if runtime_status else [],
             completedTimeFrom=helpers.new_timestamp(completed_time_from) if completed_time_from else None,
             completedTimeTo=helpers.new_timestamp(completed_time_to) if completed_time_to else None,
             pageSize=page_size or 0,
@@ -569,7 +569,7 @@ async def list_instance_ids(self,
                                 page_size: Optional[int] = None,
                                 continuation_token: Optional[str] = None) -> Page[str]:
         req = pb.ListInstanceIdsRequest(
-            runtimeStatus=[status.value for status in runtime_status] if runtime_status else None,
+            runtimeStatus=[status.value for status in runtime_status] if runtime_status else [],
             completedTimeFrom=helpers.new_timestamp(completed_time_from) if completed_time_from else None,
             completedTimeTo=helpers.new_timestamp(completed_time_to) if completed_time_to else None,
             pageSize=page_size or 0,

durabletask/history.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 from dataclasses import asdict, dataclass
-from datetime import datetime
+from datetime import datetime, timezone
 from typing import Any, Optional
 
 from google.protobuf import json_format
@@ -254,7 +254,7 @@ def to_dict(event: HistoryEvent) -> dict[str, Any]:
 def _base_kwargs(event: pb.HistoryEvent) -> dict[str, Any]:
     return {
         'event_id': event.eventId,
-        'timestamp': event.timestamp.ToDatetime(),
+        'timestamp': event.timestamp.ToDatetime(timezone.utc),
     }
 
 
@@ -266,7 +266,7 @@ def _string_value(msg: Message, field_name: str) -> Optional[str]:
 
 def _timestamp_value(msg: Message, field_name: str) -> Optional[datetime]:
     if msg.HasField(field_name):
-        return getattr(msg, field_name).ToDatetime()
+        return getattr(msg, field_name).ToDatetime(timezone.utc)
     return None
 
 
@@ -398,11 +398,11 @@ def _to_serializable(value: Any) -> Any:
     ),
     'timerCreated': lambda event: TimerCreatedEvent(
         **_base_kwargs(event),
-        fire_at=event.timerCreated.fireAt.ToDatetime(),
+        fire_at=event.timerCreated.fireAt.ToDatetime(timezone.utc),
     ),
     'timerFired': lambda event: TimerFiredEvent(
         **_base_kwargs(event),
-        fire_at=event.timerFired.fireAt.ToDatetime(),
+        fire_at=event.timerFired.fireAt.ToDatetime(timezone.utc),
         timer_id=event.timerFired.timerId,
     ),
     'orchestratorStarted': lambda event: OrchestratorStartedEvent(**_base_kwargs(event)),

durabletask/testing/in_memory_backend.py

@@ -10,6 +10,7 @@
 or external storage is not desired.
 """
 
+import bisect
 import logging
 import threading
 import time
@@ -98,6 +99,10 @@ class StateWaiter:
     result: Optional[OrchestrationInstance] = None
 
 
+_DEFAULT_PAGE_SIZE = 100
+_TOKEN_SEP = '|'
+
+
 class InMemoryOrchestrationBackend(stubs.TaskHubSidecarServiceServicer):
     """
     In-memory backend for durable orchestrations suitable for testing.
@@ -473,19 +478,24 @@ def ListInstanceIds(self, request: pb.ListInstanceIdsRequest, context):
                 matching.append(instance)
 
             matching.sort(key=lambda i: (i.completed_at, i.instance_id))
+            sort_keys = [(i.completed_at, i.instance_id) for i in matching]
 
             start_index = 0
             if request.HasField("lastInstanceKey") and request.lastInstanceKey.value:
-                for idx, instance in enumerate(matching):
-                    if instance.instance_id == request.lastInstanceKey.value:
-                        start_index = idx + 1
-                        break
-
-            page_size = request.pageSize if request.pageSize > 0 else len(matching)
+                token = request.lastInstanceKey.value
+                sep_idx = token.index(_TOKEN_SEP)
+                token_ts = datetime.fromisoformat(token[:sep_idx]).replace(tzinfo=timezone.utc)
+                token_id = token[sep_idx + 1:]
+                # bisect_right positions us just after the cursor entry
+                start_index = bisect.bisect_right(sort_keys, (token_ts, token_id))
+
+            page_size = request.pageSize if request.pageSize > 0 else _DEFAULT_PAGE_SIZE
             page = matching[start_index:start_index + page_size]
             next_token = None
             if start_index + page_size < len(matching) and page:
-                next_token = wrappers_pb2.StringValue(value=page[-1].instance_id)
+                last = page[-1]
+                encoded = f"{last.completed_at.isoformat()}{_TOKEN_SEP}{last.instance_id}"
+                next_token = wrappers_pb2.StringValue(value=encoded)
 
         return pb.ListInstanceIdsResponse(
             instanceIds=[instance.instance_id for instance in page],

tests/durabletask/test_batch_actions.py

@@ -340,7 +340,8 @@ def test_list_instance_ids_paginates_terminal_instances(backend):
     assert set(first_page.items + second_page.items) == {completed_id, failed_id}
     assert failed_state is not None
     assert failed_state.runtime_status == client.OrchestrationStatus.FAILED
-    assert first_page.continuation_token in {completed_id, failed_id}
+    assert first_page.continuation_token is not None
+    assert any(instance_id in first_page.continuation_token for instance_id in {completed_id, failed_id})
     assert second_page.continuation_token is None
 
 

tests/durabletask/test_orchestration_e2e.py

@@ -144,9 +144,12 @@ def simple(ctx: task.OrchestrationContext, value: int):
         w.start()
 
         task_hub_client = client.TaskHubGrpcClient(host_address=HOST)
-        instance_id = task_hub_client.schedule_new_orchestration(simple, input=1)
-        state = task_hub_client.wait_for_orchestration_completion(instance_id, timeout=30)
-        events = task_hub_client.get_orchestration_history(instance_id)
+        try:
+            instance_id = task_hub_client.schedule_new_orchestration(simple, input=1)
+            state = task_hub_client.wait_for_orchestration_completion(instance_id, timeout=30)
+            events = task_hub_client.get_orchestration_history(instance_id)
+        finally:
+            task_hub_client.close()
 
     assert state is not None
     assert len(events) > 0