[Brook] Added MemoryStreamPlugin for tests with no Redis.

temporalio/bridge/src/extension.rs

@@ -82,6 +82,12 @@ pub fn new_worker_extension(
             temporalio_streams::RedisStreamExtension::from_config_bytes(&config, meter)
                 .map_err(|err| PyValueError::new_err(err.to_string()))?
         }
+        #[cfg(feature = "streams")]
+        temporalio_streams::MEMORY_EXTENSION_NAME => {
+            let meter = runtime.core.telemetry().get_temporal_metric_meter();
+            temporalio_streams::MemoryStreamExtension::from_config_bytes(&config, meter)
+                .map_err(|err| PyValueError::new_err(err.to_string()))?
+        }
         #[cfg(debug_assertions)]
         test_extension::NAME => Arc::new(test_extension::TestExtension::new(&config)?),
         _ => {

temporalio/contrib/streams/redis.py

@@ -45,66 +45,27 @@
 _KEY_HEADER = "__temporal_stream_key"
 
 
-class RedisStreamPlugin(SimplePlugin):
-    """Carries Stream calls to Redis, through the temporalio-streams runtime in the bridge.
-
-    Put it on the client. A Worker takes its client's plugins, so one plugin serves Workflow
-    code, Activities and the client.
+class StreamPlugin(SimplePlugin):
+    """Base of the plugins that carry Stream calls to a temporalio-streams extension.
 
     .. warning::
         This API is experimental.
     """
 
     def __init__(
         self,
-        url: str | Sequence[str],
-        *,
-        cluster: bool = False,
-        key_prefix: str | None = None,
-        retention: timedelta | None = None,
-        safety_ttl: timedelta | None = None,
-        durable_writes: int = 0,
-        durable_timeout: timedelta | None = None,
-        response_timeout: timedelta | None = None,
-        runtime: temporalio.runtime.Runtime | None = None,
+        name: str,
+        extension_name: str,
+        config: ext_pb.RedisStreamConfig,
+        runtime: temporalio.runtime.Runtime | None,
     ) -> None:
-        """Create the plugin. It connects to Redis on the first Stream call.
-
-        Args:
-            url: A Redis URL, or the seed URLs of a cluster. Credentials and TLS come from the
-                URL.
-            cluster: Whether the URLs are of a Redis cluster.
-            key_prefix: The prefix of every key. Defaults to ``temporal-streams``.
-            retention: The default Retention of a new Stream. Defaults to 7 days.
-            safety_ttl: How long an open Workflow-owned Stream stays after its last write.
-                Defaults to 90 days.
-            durable_writes: When above zero, each Workflow flush waits for this many replicas.
-            durable_timeout: How long that wait takes at most. Defaults to 1 s.
-            response_timeout: How long one Redis call takes at most. Defaults to 5 s.
-            runtime: The runtime of the client. Defaults to the default runtime.
-        """
-        config = ext_pb.RedisStreamConfig(
-            urls=[url] if isinstance(url, str) else list(url),
-            cluster=cluster,
-            key_prefix=key_prefix or "",
-            durable_writes=durable_writes,
-        )
-        for name, value in (
-            ("retention", retention),
-            ("safety_ttl", safety_ttl),
-            ("durable_timeout", durable_timeout),
-            ("response_timeout", response_timeout),
-        ):
-            if value is not None:
-                getattr(config, name).FromTimedelta(value)
+        """Create the plugin over the bridge's extension with ``extension_name``."""
         self._extension = temporalio.worker.WorkerExtension(
-            EXTENSION_NAME, config.SerializeToString(), runtime=runtime
+            extension_name, config.SerializeToString(), runtime=runtime
         )
         self._temporal_stream_backend = StreamBackend(self._extension)
         self._converter = temporalio.converter.DataConverter.default
-        super().__init__(
-            "temporalio.RedisStreamPlugin", interceptors=[_StreamInterceptor(self)]
-        )
+        super().__init__(name, interceptors=[_StreamInterceptor(self)])
 
     def configure_worker(
         self, config: temporalio.worker.WorkerConfig
@@ -190,6 +151,63 @@ async def on_workflow_activation_complete(
         self._temporal_stream_backend.call("BufferWorkflowCalls", request)
 
 
+class RedisStreamPlugin(StreamPlugin):
+    """Carries Stream calls to Redis, through the temporalio-streams runtime in the bridge.
+
+    Put it on the client. A Worker takes its client's plugins, so one plugin serves Workflow
+    code, Activities and the client.
+
+    .. warning::
+        This API is experimental.
+    """
+
+    def __init__(
+        self,
+        url: str | Sequence[str],
+        *,
+        cluster: bool = False,
+        key_prefix: str | None = None,
+        retention: timedelta | None = None,
+        safety_ttl: timedelta | None = None,
+        durable_writes: int = 0,
+        durable_timeout: timedelta | None = None,
+        response_timeout: timedelta | None = None,
+        runtime: temporalio.runtime.Runtime | None = None,
+    ) -> None:
+        """Create the plugin. It connects to Redis on the first Stream call.
+
+        Args:
+            url: A Redis URL, or the seed URLs of a cluster. Credentials and TLS come from the
+                URL.
+            cluster: Whether the URLs are of a Redis cluster.
+            key_prefix: The prefix of every key. Defaults to ``temporal-streams``.
+            retention: The default Retention of a new Stream. Defaults to 7 days.
+            safety_ttl: How long an open Workflow-owned Stream stays after its last write.
+                Defaults to 90 days.
+            durable_writes: When above zero, each Workflow flush waits for this many replicas.
+            durable_timeout: How long that wait takes at most. Defaults to 1 s.
+            response_timeout: How long one Redis call takes at most. Defaults to 5 s.
+            runtime: The runtime of the client. Defaults to the default runtime.
+        """
+        config = ext_pb.RedisStreamConfig(
+            urls=[url] if isinstance(url, str) else list(url),
+            cluster=cluster,
+            key_prefix=key_prefix or "",
+            durable_writes=durable_writes,
+        )
+        for name, value in (
+            ("retention", retention),
+            ("safety_ttl", safety_ttl),
+            ("durable_timeout", durable_timeout),
+            ("response_timeout", response_timeout),
+        ):
+            if value is not None:
+                getattr(config, name).FromTimedelta(value)
+        super().__init__(
+            "temporalio.RedisStreamPlugin", EXTENSION_NAME, config, runtime
+        )
+
+
 def _withdraw_run(info: temporalio.workflow.Info) -> str:
     # A retry run withdraws the Records of the run it retries.
     if info.attempt > 1 and info.continued_run_id:
@@ -285,7 +303,7 @@ def operation_token(self) -> str | None:
 
 
 class _StreamInterceptor(temporalio.client.Interceptor, temporalio.worker.Interceptor):
-    def __init__(self, plugin: RedisStreamPlugin) -> None:
+    def __init__(self, plugin: StreamPlugin) -> None:
         super().__init__()
         self._plugin = plugin
 
@@ -515,7 +533,7 @@ class _ActivityInbound(temporalio.worker.ActivityInboundInterceptor):
     def __init__(
         self,
         next: temporalio.worker.ActivityInboundInterceptor,
-        plugin: RedisStreamPlugin,
+        plugin: StreamPlugin,
     ) -> None:
         super().__init__(next)
         self._plugin = plugin
@@ -554,7 +572,7 @@ async def _close_quietly(streams: _activity.ActivityStreams) -> None:
 
 
 def _activity_streams(
-    input: temporalio.worker.ExecuteActivityInput, plugin: RedisStreamPlugin
+    input: temporalio.worker.ExecuteActivityInput, plugin: StreamPlugin
 ) -> _activity.ActivityStreams:
     # A sync Activity has no client, so the plugin gives the backend and the converter.
     info = temporalio.activity.info()

temporalio/contrib/streams/testing.py

@@ -0,0 +1,50 @@
+"""A Stream plugin for tests, with no Redis.
+
+.. warning::
+    This module is experimental.
+"""
+
+from __future__ import annotations
+
+from datetime import timedelta
+
+import temporalio.runtime
+from temporalio.contrib.streams.proto import v1 as ext_pb
+from temporalio.contrib.streams.redis import StreamPlugin
+
+MEMORY_EXTENSION_NAME = "temporalio-streams-memory"
+
+
+class MemoryStreamPlugin(StreamPlugin):
+    """Keeps Streams in this process, with the same rules as on Redis, for tests.
+
+    Put it on the client, as with ``RedisStreamPlugin``. Every Worker and client that shares one
+    plugin shares its Streams, and they go when the process ends.
+
+    .. warning::
+        This API is experimental.
+    """
+
+    def __init__(
+        self,
+        *,
+        retention: timedelta | None = None,
+        safety_ttl: timedelta | None = None,
+        runtime: temporalio.runtime.Runtime | None = None,
+    ) -> None:
+        """Create the plugin.
+
+        Args:
+            retention: The default Retention of a new Stream. Defaults to 7 days.
+            safety_ttl: How long an open Workflow-owned Stream stays after its last write.
+                Defaults to 90 days.
+            runtime: The runtime of the client. Defaults to the default runtime.
+        """
+        config = ext_pb.RedisStreamConfig()
+        if retention is not None:
+            config.retention.FromTimedelta(retention)
+        if safety_ttl is not None:
+            config.safety_ttl.FromTimedelta(safety_ttl)
+        super().__init__(
+            "temporalio.MemoryStreamPlugin", MEMORY_EXTENSION_NAME, config, runtime
+        )

tests/contrib/streams/test_redis.py

@@ -18,7 +18,12 @@
 from temporalio.client import Client
 from temporalio.common import RawValue, RetryPolicy
 from temporalio.contrib import streams
-from temporalio.contrib.streams.redis import RedisStreamPlugin, _WorkflowInbound
+from temporalio.contrib.streams.redis import (
+    RedisStreamPlugin,
+    StreamPlugin,
+    _WorkflowInbound,
+)
+from temporalio.contrib.streams.testing import MemoryStreamPlugin
 from temporalio.converter import DataConverter, PayloadCodec
 from temporalio.exceptions import ApplicationError
 from temporalio.worker import (
@@ -71,10 +76,20 @@ def redis_url() -> Iterator[str]:
         process.wait()
 
 
-@pytest.fixture
-async def stream_client(client: Client, redis_url: str) -> AsyncIterator[Client]:
+@pytest.fixture(params=["memory", "redis"])
+async def stream_client(
+    client: Client, request: pytest.FixtureRequest
+) -> AsyncIterator[Client]:
+    # Every live test runs on the memory Store and on Redis, since both must keep one set of
+    # rules.
+    plugin: StreamPlugin
+    if request.param == "memory":
+        plugin = MemoryStreamPlugin()
+    else:
+        redis_url = request.getfixturevalue("redis_url")
+        plugin = RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}")
     config = client.config()
-    config["plugins"] = [RedisStreamPlugin(redis_url, key_prefix=f"t-{uuid.uuid4()}")]
+    config["plugins"] = [plugin]
     yield Client(**config)
 
 
@@ -665,9 +680,7 @@ async def test_plugin_replays_a_history_made_without_it(
     replayer = Replayer(
         workflows=[RandomAroundActivity],
         plugins=[
-            p
-            for p in stream_client.config()["plugins"]
-            if isinstance(p, RedisStreamPlugin)
+            p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin)
         ],
     )
     await replayer.replay_workflow(history)
@@ -697,7 +710,7 @@ async def test_worker_shutdown_drains_closes(stream_client: Client):
 
 async def test_plugin_shutdown_flushes_client_writers(stream_client: Client):
     [plugin] = [
-        p for p in stream_client.config()["plugins"] if isinstance(p, RedisStreamPlugin)
+        p for p in stream_client.config()["plugins"] if isinstance(p, StreamPlugin)
     ]
     stream = await streams.create_stream(stream_client, f"shutdown-{uuid.uuid4()}")
     stream.topic("t", type=str).publish("unflushed")