Enforce at registration time

Author: Sushisource

temporalio/worker/_replayer.py

@@ -216,6 +216,7 @@ def on_eviction_hook(
                 on_eviction_hook=on_eviction_hook,
                 disable_eager_activity_execution=False,
                 disable_safe_eviction=self._config["disable_safe_workflow_eviction"],
+                should_enforce_versioning_behavior=False,
             )
             # Create bridge worker
             bridge_worker, pusher = temporalio.bridge.worker.Worker.for_replay(

temporalio/worker/_worker.py

@@ -18,8 +18,6 @@
     Optional,
     Sequence,
     Type,
-    TypeAlias,
-    Union,
     cast,
 )
 
@@ -331,6 +329,12 @@ def __init__(
             )
         self._workflow_worker: Optional[_WorkflowWorker] = None
         if workflows:
+            should_enforce_versioning_behavior = (
+                deployment_options is not None
+                and deployment_options.use_worker_versioning
+                and deployment_options.default_versioning_behavior
+                == temporalio.common.VersioningBehavior.UNSPECIFIED
+            )
             self._workflow_worker = _WorkflowWorker(
                 bridge_worker=lambda: self._bridge_worker,
                 namespace=client.namespace,
@@ -348,6 +352,7 @@ def __init__(
                 metric_meter=self._runtime.metric_meter,
                 on_eviction_hook=None,
                 disable_safe_eviction=disable_safe_workflow_eviction,
+                should_enforce_versioning_behavior=should_enforce_versioning_behavior,
             )
 
         if tuner is not None:

temporalio/worker/_workflow.py

@@ -8,7 +8,6 @@
 import os
 import sys
 import threading
-from dataclasses import dataclass
 from datetime import timezone
 from types import TracebackType
 from typing import (
@@ -78,6 +77,7 @@ def __init__(
             ]
         ],
         disable_safe_eviction: bool,
+        should_enforce_versioning_behavior: bool,
     ) -> None:
         self._bridge_worker = bridge_worker
         self._namespace = namespace
@@ -135,6 +135,25 @@ def __init__(
             # Confirm name unique
             if defn.name in self._workflows:
                 raise ValueError(f"More than one workflow named {defn.name}")
+            if should_enforce_versioning_behavior:
+                not_in_annotation = defn.versioning_behavior in [
+                    None,
+                    temporalio.common.VersioningBehavior.UNSPECIFIED,
+                ]
+                if defn.name:
+                    if not_in_annotation:
+                        raise ValueError(
+                            f"Workflow {defn.name} must specify a versioning behavior using "
+                            "the `versioning_behavior` argument to `@workflow.run`."
+                        )
+                else:
+                    if defn.dynamic_versioning_behavior is None:
+                        raise ValueError(
+                            f"Dynamic Workflow {defn.name} must specify a versioning behavior "
+                            "using `@workflow.dynamic_versioning_behavior` or the "
+                            "`versioning_behavior` argument to `@workflow.run`."
+                        )
+
             # Prepare the workflow with the runner (this will error in the
             # sandbox if an import fails somehow)
             try:

tests/test_workflow.py

@@ -430,9 +430,9 @@ def test_parameters_identical_up_to_naming():
     for f1, f2 in itertools.combinations(fns, 2):
         name1, name2 = f1.__name__, f2.__name__
         expect_equal = name1[0] == name2[0]
-        assert workflow._parameters_identical_up_to_naming(f1, f2) == (expect_equal), (
-            f"expected {name1} and {name2} parameters{' ' if expect_equal else ' not '}to compare equal"
-        )
+        assert (
+            workflow._parameters_identical_up_to_naming(f1, f2) == (expect_equal)
+        ), f"expected {name1} and {name2} parameters{' ' if expect_equal else ' not '}to compare equal"
 
 
 @workflow.defn

tests/worker/test_worker.py

@@ -843,7 +843,94 @@ async def test_worker_deployment_dynamic_workflow_getter(
     )
 
 
-# TODO: Test for fail at registration time if deployment versioning on, no default, no behavior
+@workflow.defn
+class NoVersioningAnnotationWorkflow:
+    @workflow.run
+    async def run(self) -> str:
+        return "whee"
+
+
+@workflow.defn(dynamic=True)
+class NoVersioningAnnotationDynamicWorkflow:
+    @workflow.run
+    async def run(self, args: Sequence[RawValue]) -> str:
+        return "whee"
+
+
+async def test_workflows_must_have_versioning_behavior_when_feature_turned_on(
+    client: Client, env: WorkflowEnvironment
+):
+    with pytest.raises(ValueError) as exc_info:
+        Worker(
+            client,
+            task_queue=f"task-queue-{uuid.uuid4()}",
+            workflows=[NoVersioningAnnotationWorkflow],
+            deployment_options=WorkerDeploymentOptions(
+                version=WorkerDeploymentVersion(
+                    deployment_name="whatever", build_id="1.0"
+                ),
+                use_worker_versioning=True,
+            ),
+        )
+
+    assert "must specify a versioning behavior" in str(exc_info.value)
+
+    with pytest.raises(ValueError) as exc_info:
+        Worker(
+            client,
+            task_queue=f"task-queue-{uuid.uuid4()}",
+            workflows=[NoVersioningAnnotationDynamicWorkflow],
+            deployment_options=WorkerDeploymentOptions(
+                version=WorkerDeploymentVersion(
+                    deployment_name="whatever", build_id="1.0"
+                ),
+                use_worker_versioning=True,
+            ),
+        )
+
+    assert "must specify a versioning behavior" in str(exc_info.value)
+
+
+async def test_workflows_can_use_default_versioning_behavior(
+    client: Client, env: WorkflowEnvironment
+):
+    if env.supports_time_skipping:
+        pytest.skip("Test Server doesn't support worker versioning")
+
+    deployment_name = f"deployment-default-versioning-{uuid.uuid4()}"
+    worker_v1 = WorkerDeploymentVersion(deployment_name=deployment_name, build_id="1.0")
+
+    async with new_worker(
+        client,
+        NoVersioningAnnotationWorkflow,
+        deployment_options=WorkerDeploymentOptions(
+            version=worker_v1,
+            use_worker_versioning=True,
+            default_versioning_behavior=VersioningBehavior.PINNED,
+        ),
+    ) as w:
+        describe_resp = await wait_until_worker_deployment_visible(
+            client,
+            worker_v1,
+        )
+        await set_current_deployment_version(
+            client, describe_resp.conflict_token, worker_v1
+        )
+
+        wf = await client.start_workflow(
+            NoVersioningAnnotationWorkflow.run,
+            id=f"default-versioning-behavior-{uuid.uuid4()}",
+            task_queue=w.task_queue,
+        )
+        await wf.result()
+
+        history = await wf.fetch_history()
+        assert any(
+            event.HasField("workflow_task_completed_event_attributes")
+            and event.workflow_task_completed_event_attributes.versioning_behavior
+            == temporalio.api.enums.v1.VersioningBehavior.VERSIONING_BEHAVIOR_PINNED
+            for event in history.events
+        )
 
 
 async def wait_until_worker_deployment_visible(