https://community.databricks.com/t5/technical-blog/unit-testing-transformwithstate-logic-with-twstester/ba-p/156085 <P>You've spent the afternoon building a <CODE>StatefulProcessor</CODE> for your TransformWithState streaming job. It tracks per-user sessions, accumulates running totals, or deduplicates events. Now you want to know if it actually works.</P> <P>So you wire up a streaming source, start a query, push some test rows, wait for a micro-batch to fire, check the sink, squint at the logs, tweak something, restart, and do it all again. The feedback loop is slow, the state is invisible, and reproducing a specific sequence of inputs feels like guesswork.</P> <P>TwsTester changes that. It's a small test helper, available in both PySpark and Scala starting in Spark 4.2.0 (DBR 18.2). You hand it your <CODE>StatefulProcessor</CODE>, feed it rows, and inspect outputs and internal state instantly, without standing up a streaming query.</P> <H2 id="what-is-twstester">What is TwsTester?</H2> <P>TwsTester is a unit-testing harness for <CODE>StatefulProcessor</CODE> implementations used with the <CODE>transformWithState</CODE> operator in Structured Streaming. You hand it your processor and some input rows; it drives the processor the same way Spark would during a real micro-batch and returns the resulting output rows.</P> <P>It's available at:</P> <UL> <LI><STRONG>PySpark</STRONG>: <CODE>pyspark.sql.streaming.TwsTester</CODE></LI> <LI><STRONG>Scala</STRONG>: <CODE>org.apache.spark.sql.streaming.TwsTester</CODE></LI> </UL> <P>The emphasis is on small, focused tests you can run in a notebook cell or a CI pipeline, not full end-to-end streaming tests with sources and sinks. Think of it as the difference between calling a function in a unit test and spinning up a full integration environment to exercise the same code path.</P> <H2 id="why-twstester">Why TwsTester?</H2> <P>Three pain points come up repeatedly when developing stateful streaming logic:</P> <OL type="1"> <LI><STRONG>Slow feedback loops.</STRONG> Every change to your processor means restarting a streaming query, wiring up sources and sinks, and waiting for data to flow through. For logic that's still taking shape, this is expensive.</LI> <LI><STRONG>Hard-to-observe state and timers.</STRONG> In a running query, inspecting state mid-flight isn't straightforward. The <A href="https://www.databricks.com/blog/announcing-state-reader-api-new-statestore-data-source" target="_blank" rel="noopener">State Reader API</A> helps here; it lets you read state from a query's checkpoint directory after the fact, which is valuable for debugging production workloads and understanding <A href="https://www.databricks.com/blog/announcing-simplified-state-tracking-apache-sparktm-structured-streaming" target="_blank" rel="noopener">state evolution over time</A>. But the State Reader API requires access to the checkpoint directory, and it's oriented toward inspecting state <EM>after</EM> a query has run, not toward feeding controlled inputs and verifying outputs as part of a development loop.</LI> <LI><STRONG>Limited reproducibility.</STRONG> Some bugs only surface with specific input orderings or timing. Reproducing those conditions end-to-end is tedious at best.</LI> </OL> <P>TwsTester addresses all three:</P> <UL> <LI>Process a <STRONG>controlled sequence of input rows</STRONG> per key and inspect the exact outputs.</LI> <LI><STRONG>Peek at and manipulate state</STRONG> directly, including value state, list state, and map state, at any point during a test.</LI> <LI><STRONG>Advance processing time or the watermark</STRONG> on demand, and verify that your timer logic fires correctly.</LI> <LI><STRONG>Inject initial state</STRONG> so you can test mid-stream scenarios like "this user already has an active session."</LI> </UL> <P>This all runs locally against a regular <CODE>SparkSession</CODE>. No cluster infrastructure, no streaming sources, no sinks.</P> <H2 id="how-it-works">How it works</H2> <P>The workflow is straightforward:</P> <OL type="1"> <LI><STRONG>Implement your StatefulProcessor</STRONG> as you normally would for <CODE>transformWithState</CODE>: define <CODE>init</CODE>, <CODE>handleInputRows</CODE>, and optionally <CODE>handleExpiredTimer</CODE> and <CODE>close</CODE>.</LI> <LI><STRONG>Construct a TwsTester</STRONG> with your processor, choosing a time mode (<CODE>"None"</CODE>, <CODE>"ProcessingTime"</CODE>, or <CODE>"EventTime"</CODE>) and output mode (<CODE>"Append"</CODE>, <CODE>"Update"</CODE>, or <CODE>"Complete"</CODE>).</LI> <LI><STRONG>Call <CODE>test(key, rows)</CODE></STRONG> with a grouping key and a list of input rows. TwsTester drives your processor and returns the output rows.</LI> <LI><STRONG>Inspect state</STRONG> with <CODE>peekValueState</CODE>, <CODE>peekListState</CODE>, or <CODE>peekMapState</CODE>. Seed or override state with the corresponding <CODE>update*</CODE> methods.</LI> <LI><STRONG>Advance time</STRONG> with <CODE>setProcessingTime</CODE> or <CODE>setWatermark</CODE> to fire expired timers and collect their output.</LI> </OL> <P>Steps 3-5 can be repeated in any order, which makes it easy to build up multi-step test scenarios.</P> <H2 id="example-catching-a-bug-with-twstester">Example: catching a bug with TwsTester</H2> <P>Let's walk through a realistic development scenario. We'll build a processor that maintains a running sum of transaction amounts per user, write tests for it, discover a bug, fix it, and verify the fix. This is exactly the kind of tight iteration loop that TwsTester enables.</P> <H3 id="the-processor-first-attempt">The processor (first attempt)</H3> <P>Here's our first pass at a <CODE>RunningSumProcessor</CODE>. It takes input rows with an <CODE>amount</CODE> field, accumulates a per-user total, and emits the updated running sum with each row.</P> <LI-CODE lang="python">from pyspark.sql import Row, SparkSession from pyspark.sql.streaming import StatefulProcessor, StatefulProcessorHandle from pyspark.sql.streaming import TwsTester from pyspark.sql.types import StructType, StructField, LongType spark = SparkSession.builder.getOrCreate() class RunningSumProcessor(StatefulProcessor): def init(self, handle: StatefulProcessorHandle) -&gt; None: self.handle = handle sum_schema = StructType([StructField("sum", LongType())]) self.sum_state = handle.getValueState("sum", sum_schema) def handleInputRows(self, key, rows, timerValues): running_sum = 0 results = [] for row in rows: amount = row["amount"] running_sum += amount results.append(Row(user_id=key[0], amount=amount, running_sum=running_sum)) self.sum_state.update((running_sum,)) return iter(results) def close(self) -&gt; None: pass</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">import org.apache.spark.sql.streaming._ import org.apache.spark.sql.Encoders class RunningSumProcessor extends StatefulProcessor[String, Long, (String, Long, Long)] { @transient private var sumState: ValueState[Long] = _ override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { sumState = getHandle.getValueState[Long]( "sum", Encoders.scalaLong, TTLConfig.NONE) } override def handleInputRows( key: String, rows: Iterator[Long], timerValues: TimerValues ): Iterator[(String, Long, Long)] = { var runningSum = 0L val results = rows.map { amount =&gt; runningSum += amount (key, amount, runningSum) }.toList sumState.update(runningSum) results.iterator } }</LI-CODE></DETAILS> <P>Looks reasonable. Let's test it.</P> <H3 id="setting-up-twstester">Setting up TwsTester</H3> <LI-CODE lang="python">tester = TwsTester( processor=RunningSumProcessor(), timeMode="None", outputMode="Append", )</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">val tester = new TwsTester( processor = new RunningSumProcessor(), timeMode = TimeMode.None(), outputMode = OutputMode.Append() )</LI-CODE></DETAILS> <P>No sources, no sinks, no checkpoint directory.</P> <H3 id="first-test-single-batch">First test: single batch</H3> <P>We process two transactions for Alice in one call. The running sum should accumulate within the batch.</P> <LI-CODE lang="python">output = tester.test( key="alice", input=[Row(amount=10), Row(amount=5)], ) for row in output: print(row) # Row(user_id='alice', amount=10, running_sum=10) # Row(user_id='alice', amount=5, running_sum=15)</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">val output = tester.test("alice", List(10L, 5L)) println(output) // List(("alice", 10, 10), ("alice", 5, 15))</LI-CODE></DETAILS> <P>Looks right. Alice's running sum goes 10, then 15.</P> <H3 id="second-test-another-batch-for-the-same-key">Second test: another batch for the same key</H3> <P>Now Alice makes another transaction. She already has a running sum of 15, so the new total should be 18.</P> <LI-CODE lang="python">output2 = tester.test( key="alice", input=[Row(amount=3)], ) for row in output2: print(row) # Expected: Row(user_id='alice', amount=3, running_sum=18) # Actual: Row(user_id='alice', amount=3, running_sum=3) &lt;- wrong!</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">val output2 = tester.test("alice", List(3L)) println(output2) // Expected: List(("alice", 3, 18)) // Actual: List(("alice", 3, 3)) &lt;- wrong!</LI-CODE></DETAILS> <P>The running sum is 3, not 18. The processor reset Alice's total instead of continuing from where it left off.</P> <H3 id="what-went-wrong">What went wrong</H3> <P>Look at the first line of <CODE>handleInputRows</CODE>:</P> <LI-CODE lang="python">running_sum = 0 # Always starts from zero, ignoring existing state</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">var runningSum = 0L // Always starts from zero, ignoring existing state</LI-CODE></DETAILS> <P>The processor initializes <CODE>running_sum</CODE> to zero on every call instead of reading the existing value from state. Within a single batch this doesn't matter because the accumulation happens in a loop. But across batches, the state from previous calls is silently discarded.</P> <P>Without TwsTester, this bug might not surface until production, where it would show up as totals that periodically "reset" at unpredictable intervals (whenever a new micro-batch starts for that key).</P> <H3 id="fixing-the-processor">Fixing the processor</H3> <P>The fix is two lines: read existing state before accumulating.</P> <LI-CODE lang="python">class RunningSumProcessor(StatefulProcessor): def init(self, handle: StatefulProcessorHandle) -&gt; None: self.handle = handle sum_schema = StructType([StructField("sum", LongType())]) self.sum_state = handle.getValueState("sum", sum_schema) def handleInputRows(self, key, rows, timerValues): # Fixed: read existing state instead of starting from zero existing = self.sum_state.get() running_sum = existing[0] if existing is not None else 0 results = [] for row in rows: amount = row["amount"] running_sum += amount results.append(Row(user_id=key[0], amount=amount, running_sum=running_sum)) self.sum_state.update((running_sum,)) return iter(results) def close(self) -&gt; None: pass</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">class RunningSumProcessor extends StatefulProcessor[String, Long, (String, Long, Long)] { @transient private var sumState: ValueState[Long] = _ override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { sumState = getHandle.getValueState[Long]( "sum", Encoders.scalaLong, TTLConfig.NONE) } override def handleInputRows( key: String, rows: Iterator[Long], timerValues: TimerValues ): Iterator[(String, Long, Long)] = { // Fixed: read existing state instead of starting from zero var runningSum = if (sumState.exists()) sumState.get() else 0L val results = rows.map { amount =&gt; runningSum += amount (key, amount, runningSum) }.toList sumState.update(runningSum) results.iterator } }</LI-CODE></DETAILS> <H3 id="re-running-the-tests">Re-running the tests</H3> <LI-CODE lang="python">tester = TwsTester( processor=RunningSumProcessor(), # fixed version timeMode="None", outputMode="Append", ) # Test 1: single batch output = tester.test(key="alice", input=[Row(amount=10), Row(amount=5)]) for row in output: print(row) # Row(user_id='alice', amount=10, running_sum=10) # Row(user_id='alice', amount=5, running_sum=15) # Test 2: second batch for the same key output2 = tester.test(key="alice", input=[Row(amount=3)]) for row in output2: print(row) # Row(user_id='alice', amount=3, running_sum=18) PASS</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">val tester = new TwsTester( processor = new RunningSumProcessor(), // fixed version timeMode = TimeMode.None(), outputMode = OutputMode.Append() ) // Test 1: single batch val output = tester.test("alice", List(10L, 5L)) println(output) // List(("alice", 10, 10), ("alice", 5, 15)) // Test 2: second batch for the same key val output2 = tester.test("alice", List(3L)) println(output2) // List(("alice", 3, 18)) PASS</LI-CODE></DETAILS> <P>Both tests pass. From bug to fix to verification in seconds, not the minutes (or longer) it would take to redeploy a streaming query and push test data through a source.</P> <H3 id="inspecting-state-directly">Inspecting state directly</H3> <P>After calling <CODE>test()</CODE>, the processor's state is still live inside the tester. We can peek at it:</P> <LI-CODE lang="python">alice_sum = tester.peekValueState("sum", "alice") print(alice_sum) # (18,)</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">val aliceSum = tester.peekValueState[Long]("sum", "alice") println(aliceSum) // Some(18)</LI-CODE></DETAILS> <P>We can also process rows for a different key and verify that state is isolated:</P> <LI-CODE lang="python">output_bob = tester.test(key="bob", input=[Row(amount=7)]) print(output_bob) # [Row(user_id='bob', amount=7, running_sum=7)] # Bob's state is independent print(tester.peekValueState("sum", "bob")) # (7,) print(tester.peekValueState("sum", "alice")) # (18,) -- unchanged</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">val outputBob = tester.test("bob", List(7L)) println(outputBob) // List(("bob", 7, 7)) // Bob's state is independent println(tester.peekValueState[Long]("sum", "bob")) // Some(7) println(tester.peekValueState[Long]("sum", "alice")) // Some(18) -- unchanged</LI-CODE></DETAILS> <H3 id="seeding-state-for-mid-stream-scenarios">Seeding state for mid-stream scenarios</H3> <P>Sometimes you want to test behavior starting from a known state rather than building up from zero. <CODE>updateValueState</CODE> lets you seed state before processing:</P> <LI-CODE lang="python">tester.updateValueState("sum", "carol", (100,)) output_carol = tester.test(key="carol", input=[Row(amount=25)]) print(output_carol) # [Row(user_id='carol', amount=25, running_sum=125)]</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">tester.updateValueState[Long]("sum", "carol", 100L) val outputCarol = tester.test("carol", List(25L)) println(outputCarol) // List(("carol", 25, 125))</LI-CODE></DETAILS> <P>This is useful for testing edge cases: what happens at counter rollover, after a long idle period, or when the user already has existing state from a previous session.</P> <H2 id="adding-timer-logic">Adding timer logic</H2> <P>Timers let a <CODE>StatefulProcessor</CODE> react to the passage of time: session timeouts, periodic flushes, TTL-based cleanup. TwsTester supports testing timer logic by letting you advance time manually.</P> <P>Here's a processor that registers a processing-time timer to expire idle sessions:</P> <LI-CODE lang="python">class SessionTimeoutProcessor(StatefulProcessor): TIMEOUT_MS = 60_000 # 1 minute def init(self, handle: StatefulProcessorHandle) -&gt; None: self.handle = handle schema = StructType([StructField("event_count", LongType())]) self.session = handle.getValueState("session", schema) def handleInputRows(self, key, rows, timerValues): existing = self.session.get() event_count = existing[0] if existing is not None else 0 for row in rows: event_count += 1 self.session.update((event_count,)) # Reset the timeout: register a timer at current time + timeout current_time = timerValues.getCurrentProcessingTimeInMs() self.handle.registerTimer(current_time + self.TIMEOUT_MS) return iter([Row(user_id=key[0], event_count=event_count, status="active")]) def handleExpiredTimer(self, key, timerValues, expiredTimerInfo): # Timer fired, session timed out. Clear state and emit a closure event. event_count = self.session.get() self.session.clear() count = event_count[0] if event_count is not None else 0 return iter([Row(user_id=key[0], event_count=count, status="expired")]) def close(self) -&gt; None: pass</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">class SessionTimeoutProcessor extends StatefulProcessor[String, String, (String, Long, String)] { private val TimeoutMs = 60000L @transient private var sessionState: ValueState[Long] = _ override def init(outputMode: OutputMode, timeMode: TimeMode): Unit = { sessionState = getHandle.getValueState[Long]( "session", Encoders.scalaLong, TTLConfig.NONE) } override def handleInputRows( key: String, rows: Iterator[String], timerValues: TimerValues ): Iterator[(String, Long, String)] = { var eventCount = if (sessionState.exists()) sessionState.get() else 0L rows.foreach(_ =&gt; eventCount += 1) sessionState.update(eventCount) val currentTime = timerValues.getCurrentProcessingTimeInMs() getHandle.registerTimer(currentTime + TimeoutMs) Iterator((key, eventCount, "active")) } override def handleExpiredTimer( key: String, timerValues: TimerValues, expiredTimerInfo: ExpiredTimerInfo ): Iterator[(String, Long, String)] = { val count = if (sessionState.exists()) sessionState.get() else 0L sessionState.clear() Iterator((key, count, "expired")) } }</LI-CODE></DETAILS> <P>Now test it with TwsTester in processing-time mode:</P> <LI-CODE lang="python">tester = TwsTester( processor=SessionTimeoutProcessor(), timeMode="ProcessingTime", outputMode="Append", ) # Set an initial processing time tester.setProcessingTime(1_000) # Process some events for alice output = tester.test( key="alice", input=[Row(event="click"), Row(event="scroll")], ) print(output) # [Row(user_id='alice', event_count=2, status='active')] # Advance time past the timeout threshold expired_output = tester.setProcessingTime(62_000) print(expired_output) # [Row(user_id='alice', event_count=2, status='expired')] # State has been cleared print(tester.peekValueState("session", "alice")) # None</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show Scala equivalent</SUMMARY> <LI-CODE lang="scala">val tester = new TwsTester( processor = new SessionTimeoutProcessor(), timeMode = TimeMode.ProcessingTime(), outputMode = OutputMode.Append() ) // Set an initial processing time tester.setProcessingTime(1000L) // Process some events for alice val output = tester.test("alice", List("click", "scroll")) println(output) // List(("alice", 2, "active")) // Advance time past the timeout threshold val expiredOutput = tester.setProcessingTime(62000L) println(expiredOutput) // List(("alice", 2, "expired")) // State has been cleared println(tester.peekValueState[Long]("session", "alice")) // None</LI-CODE></DETAILS> <P>We just tested a complete session lifecycle (activity, timeout, cleanup) in a few lines, with full control over time.</P> <H2 id="wrapping-tests-for-ci">Wrapping tests for CI</H2> <P>The examples above work well in a notebook for interactive exploration. To lock in behavior for CI, wrap them in your test framework of choice.</P> <P>Here's a pytest example. The same pattern works in ScalaTest; expand the spoiler below for the Scala version.</P> <LI-CODE lang="python">import pytest from pyspark.sql import Row, SparkSession from pyspark.sql.streaming import TwsTester @pytest.fixture(scope="module") def spark(): return SparkSession.builder.master("local[*]").getOrCreate() def test_running_sum_single_batch(spark): tester = TwsTester( processor=RunningSumProcessor(), timeMode="None", outputMode="Append", ) output = tester.test("alice", [Row(amount=10), Row(amount=5)]) assert len(output) == 2 assert output[0]["running_sum"] == 10 assert output[1]["running_sum"] == 15 def test_running_sum_across_batches(spark): tester = TwsTester( processor=RunningSumProcessor(), timeMode="None", outputMode="Append", ) tester.test("alice", [Row(amount=10), Row(amount=5)]) output = tester.test("alice", [Row(amount=3)]) assert output[0]["running_sum"] == 18 def test_state_isolation_across_keys(spark): tester = TwsTester( processor=RunningSumProcessor(), timeMode="None", outputMode="Append", ) tester.test("alice", [Row(amount=10)]) tester.test("bob", [Row(amount=20), Row(amount=5)]) assert tester.peekValueState("sum", "alice") == (10,) assert tester.peekValueState("sum", "bob") == (25,)</LI-CODE><DETAILS style="margin: 8px 0 16px 0; border-left: 3px solid #FF3621; padding-left: 12px;" open="open"> <SUMMARY style="cursor: pointer; color: #1f6feb; font-weight: 500; padding: 4px 0; user-select: none;">▸ Show ScalaTest equivalent</SUMMARY> <LI-CODE lang="scala">import org.scalatest.funsuite.AnyFunSuite import org.apache.spark.sql.streaming._ class RunningSumProcessorTest extends AnyFunSuite { test("running sum accumulates within a single batch") { val tester = new TwsTester( processor = new RunningSumProcessor(), timeMode = TimeMode.None(), outputMode = OutputMode.Append() ) val output = tester.test("alice", List(10L, 5L)) assert(output === List(("alice", 10L, 10L), ("alice", 5L, 15L))) } test("running sum accumulates across batches") { val tester = new TwsTester( processor = new RunningSumProcessor(), timeMode = TimeMode.None(), outputMode = OutputMode.Append() ) tester.test("alice", List(10L, 5L)) val output = tester.test("alice", List(3L)) assert(output === List(("alice", 3L, 18L))) } test("state is isolated across keys") { val tester = new TwsTester( processor = new RunningSumProcessor(), timeMode = TimeMode.None(), outputMode = OutputMode.Append() ) tester.test("alice", List(10L)) tester.test("bob", List(20L, 5L)) assert(tester.peekValueState[Long]("sum", "alice") === Some(10L)) assert(tester.peekValueState[Long]("sum", "bob") === Some(25L)) } }</LI-CODE></DETAILS> <P>No mock streaming sources, no <CODE>awaitTermination</CODE>, no checkpoint cleanup. Just direct assertions on processor behavior.</P> <H2 id="fitting-twstester-into-your-workflow">Fitting TwsTester into your workflow</H2> <P>TwsTester is most valuable when you're in the thick of building or debugging stateful logic:</P> <UL> <LI><STRONG>During development</STRONG>: run tests in a notebook cell after each change. The feedback loop drops from minutes to seconds.</LI> <LI><STRONG>For corner cases</STRONG>: construct specific input sequences that exercise edge conditions such as late arrivals, duplicate keys, empty batches, timer expirations at exact boundaries.</LI> <LI><STRONG>In CI</STRONG>: wrap the same tests in pytest or ScalaTest so regressions are caught before they reach a streaming job.</LI> </UL> <P>It's not a replacement for end-to-end integration tests. You'll still want to verify that your processor works correctly within a full <CODE>transformWithState</CODE> pipeline. But for the business logic inside the processor, TwsTester gives you the tight feedback loop that stateful streaming has always been missing.</P> <H2 id="get-started">Get started</H2> <P>Copy the example above into a Databricks Notebook or a local Spark environment running Spark 4.2.0+, replace <CODE>RunningSumProcessor</CODE> with your own processor, and start testing.</P> <P>For the full API reference and more examples:</P> <UL> <LI><A href="https://spark.apache.org/docs/latest/structured-streaming-programming-guide.html#arbitrary-stateful-operations-with-transformwithstate" target="_blank" rel="noopener">TransformWithState programming guide</A> covers <CODE>StatefulProcessor</CODE>, state variables, and timer semantics.</LI> <LI><A href="https://github.com/apache/spark/pull/53159" target="_blank" rel="noopener">Scala TwsTester PR</A>: the original implementation with comprehensive test cases.</LI> <LI><A href="https://github.com/apache/spark/pull/53758" target="_blank" rel="noopener">PySpark TwsTester PR</A>: PySpark port with Row and Pandas mode examples.</LI> <LI><A href="https://github.com/apache/spark/pull/53802" target="_blank" rel="noopener">TwsTester documentation PR</A>: the programming guide additions for TwsTester.</LI> </UL> <H2><STRONG>Frequently asked questions (FAQ)</STRONG></H2> <P><STRONG>What is TwsTester?</STRONG><BR aria-hidden="true" />TwsTester is a unit-testing harness for <CODE>StatefulProcessor</CODE> implementations used with the <CODE>transformWithState</CODE> operator in Apache Spark Structured Streaming. You hand it your processor and a list of input rows, and it drives the processor the same way Spark would inside a real micro-batch, returning the resulting output rows so you can assert on them directly.</P> <P><BR aria-hidden="true" /><STRONG>Which Spark and Databricks Runtime versions include TwsTester?</STRONG><BR aria-hidden="true" />TwsTester is available starting in <STRONG>Apache Spark 4.2.0</STRONG>, which corresponds to <STRONG>Databricks Runtime 18.2</STRONG> and later.</P> <P><BR aria-hidden="true" /><STRONG>Is TwsTester available in both PySpark and Scala?</STRONG><BR aria-hidden="true" />Yes. The PySpark entry point is <CODE>pyspark.sql.streaming.TwsTester</CODE> and the Scala entry point is <CODE>org.apache.spark.sql.streaming.TwsTester</CODE>. The two surfaces mirror each other, so the same test patterns translate directly between languages.</P> <P><BR aria-hidden="true" /><STRONG>Does TwsTester replace end-to-end streaming integration tests?</STRONG><BR aria-hidden="true" />No. TwsTester targets the business logic inside a <CODE>StatefulProcessor</CODE>: per-key state, timers, output rows for given inputs. You should still run end-to-end tests against a real <CODE>transformWithState</CODE> pipeline to verify source/sink wiring, checkpointing, and fault tolerance. TwsTester compresses the <I>inner</I> feedback loop; integration tests still validate the <I>outer</I> one.</P> <P><BR aria-hidden="true" /><STRONG>Can TwsTester test processing-time and event-time timers?</STRONG><BR aria-hidden="true" />Yes. Construct the tester with <CODE>timeMode="ProcessingTime"</CODE> or <CODE>timeMode="EventTime"</CODE> and use <CODE>setProcessingTime(...)</CODE> or <CODE>setWatermark(...)</CODE> to advance time on demand. Any expired timers fire and their output is collected, so timeout, session-expiry, and periodic-flush logic can be exercised deterministically.</P> <P><BR aria-hidden="true" /><STRONG>How is TwsTester different from the State Reader API?</STRONG><BR aria-hidden="true" />The <SPAN draggable="true"><A href="https://www.databricks.com/blog/announcing-state-reader-api-new-statestore-data-source" target="_blank" rel="noopener noreferrer">State Reader API</A></SPAN> reads state from a query's checkpoint directory after the fact, which is great for inspecting production state. TwsTester runs locally, drives the processor with controlled inputs, and lets you inspect or override state mid-test with <CODE>peekValueState</CODE> / <CODE>updateValueState</CODE> and friends. They are complementary: State Reader for production debugging, TwsTester for development and CI.</P> <P><STRONG>Can I use TwsTester in a CI pipeline?</STRONG><BR aria-hidden="true" />Yes. Wrap each test in your framework of choice (<CODE>pytest</CODE> for PySpark, <CODE>ScalaTest</CODE> for Scala) and run them like any other unit test. TwsTester needs only a regular <CODE>SparkSession</CODE>, no streaming sources, sinks, or checkpoint directory, so it fits naturally into existing test infrastructure.</P> Tue, 12 May 2026 14:00:33 GMT craig_lukasik 2026-05-12T14:00:33Z https://community.databricks.com/t5/technical-blog/unit-testing-transformwithstate-logic-with-twstester/ba-p/156085 <P>TwsTester is a unit-testing harness for StatefulProcessor implementations used with transformWithState. Hand it your processor and some input rows, and it drives them the same way Spark would in a real micro-batch — with full control over state and timers.</P> Tue, 12 May 2026 14:00:33 GMT https://community.databricks.com/t5/technical-blog/unit-testing-transformwithstate-logic-with-twstester/ba-p/156085 craig_lukasik 2026-05-12T14:00:33Z