Fix replace() to register unregistered functions

src/docket/docket.py

@@ -441,6 +441,8 @@ def replace(
         """
         if isinstance(function, str):
             function = self.tasks[function]
+        else:
+            self.register(function)
 
         async def scheduler(*args: P.args, **kwargs: P.kwargs) -> Execution:
             execution = Execution(self, function, args, kwargs, key, when, attempt=1)

tests/test_fundamentals.py

@@ -178,6 +178,24 @@ async def test_rescheduling_by_name(
     assert later <= now()
 
 
+async def test_replace_without_existing_task_acts_like_add(
+    docket: Docket, worker: Worker, the_task: AsyncMock, now: Callable[[], datetime]
+):
+    """docket.replace() on a non-existent key should schedule the task like add()"""
+
+    key = f"my-cool-task:{uuid4()}"
+
+    # Replace without prior add - should just schedule the task
+    later = now() + timedelta(milliseconds=100)
+    await docket.replace(the_task, later, key=key)("b", "c", c="d")
+
+    await worker.run_until_finished()
+
+    the_task.assert_awaited_once_with("b", "c", c="d")
+
+    assert later <= now()
+
+
 async def test_task_keys_are_idempotent_in_the_future(
     docket: Docket, worker: Worker, the_task: AsyncMock, now: Callable[[], datetime]
 ):