Add AdvanceWallClockTime support for testing PROCESSING_TIME punctuations

core/Mock/ClusterInMemoryTopologyDriver.cs

@@ -244,6 +244,16 @@ public void TriggerCommit()
             //((StreamThread)threadTopology)?.Manager.CommitAll();
         }
 
+        public void AdvanceWallClockTime(TimeSpan advance)
+        {
+            // Wall-clock time advancement is not supported in the cluster in-memory driver
+            // because it uses real Kafka with threads that use actual system time.
+            // Use TaskSynchronousTopologyDriver instead for testing PROCESSING_TIME punctuations.
+            throw new NotSupportedException(
+                "AdvanceWallClockTime is not supported in ClusterInMemoryTopologyDriver. " +
+                "Use the synchronous topology driver (default mode) instead for testing PROCESSING_TIME punctuations.");
+        }
+
         #endregion
     }
 }

core/Mock/IBehaviorTopologyTestDriver.cs

@@ -15,5 +15,12 @@ internal interface IBehaviorTopologyTestDriver : IDisposable
         TestMultiInputTopic<K, V> CreateMultiInputTopic<K, V>(string[] topics, ISerDes<K> keySerdes = null, ISerDes<V> valueSerdes = null);
         IStateStore GetStateStore<K, V>(string name);
         void TriggerCommit();
+
+        /// <summary>
+        /// Advance wall-clock time by the specified duration, triggering any PROCESSING_TIME punctuations
+        /// that are scheduled to fire during the advanced time period.
+        /// </summary>
+        /// <param name="advance">The amount of time to advance the wall clock</param>
+        void AdvanceWallClockTime(TimeSpan advance);
     }
 }

core/Mock/IWallClockTimeProvider.cs

@@ -0,0 +1,22 @@
+using System;
+
+namespace Streamiz.Kafka.Net.Mock
+{
+    /// <summary>
+    /// Interface for providing wall-clock time. Used to allow time mocking in tests.
+    /// </summary>
+    internal interface IWallClockTimeProvider
+    {
+        /// <summary>
+        /// Gets the current wall-clock time in milliseconds since epoch.
+        /// </summary>
+        long WallClockTime { get; }
+
+        /// <summary>
+        /// Advances the wall-clock time by the specified duration.
+        /// Only applicable for mock implementations; production implementations may ignore this.
+        /// </summary>
+        /// <param name="advance">The amount of time to advance</param>
+        void Advance(TimeSpan advance);
+    }
+}

core/Mock/MockWallClockTimeProvider.cs

@@ -0,0 +1,24 @@
+using System;
+using Streamiz.Kafka.Net.Crosscutting;
+
+namespace Streamiz.Kafka.Net.Mock
+{
+    /// <summary>
+    /// A mock wall-clock time provider for testing that starts at the current time
+    /// and allows manual advancement.
+    /// </summary>
+    internal class MockWallClockTimeProvider : IWallClockTimeProvider
+    {
+        public MockWallClockTimeProvider()
+        {
+            WallClockTime = DateTime.Now.GetMilliseconds();
+        }
+
+        public long WallClockTime { get; private set; }
+
+        public void Advance(TimeSpan advance)
+        {
+            WallClockTime += (long)advance.TotalMilliseconds;
+        }
+    }
+}

core/Mock/Sync/SyncPipeBuilder.cs

@@ -34,11 +34,12 @@ public void PublishRecord(string topic, byte[] key, byte[] value, DateTime times
 
             public void Flush()
             {
-                long now = DateTime.Now.GetMilliseconds();
+                // Use the task's wall clock time which is mock-aware
+                long now = task.WallClockTime;
                 TaskManager.CurrentTask = task;
                 while (task.CanProcess(now))
                     task.Process();
-                
+
                 task.PunctuateStreamTime();
                 task.PunctuateSystemTime();
                 TaskManager.CurrentTask = null;

core/Mock/Sync/SyncProducer.cs

@@ -163,6 +163,17 @@ public IEnumerable<Message<byte[], byte[]>> GetHistory(string topicName)
             return topics[topicName].ToArray();
         }
 
+        public IDictionary<string, int> GetAllTopicCounts()
+        {
+            lock (_lock)
+            {
+                var result = new Dictionary<string, int>();
+                foreach (var kvp in topics)
+                    result[kvp.Key] = kvp.Value.Count;
+                return result;
+            }
+        }
+
         public void CommitTransaction()
         {
             throw new NotImplementedException();

core/Mock/SystemWallClockTimeProvider.cs

@@ -0,0 +1,24 @@
+using System;
+using Streamiz.Kafka.Net.Crosscutting;
+
+namespace Streamiz.Kafka.Net.Mock
+{
+    /// <summary>
+    /// Default wall-clock time provider that returns the actual system time.
+    /// Used in production when no mock time provider is injected.
+    /// </summary>
+    internal class SystemWallClockTimeProvider : IWallClockTimeProvider
+    {
+        /// <summary>
+        /// Singleton instance for shared use.
+        /// </summary>
+        public static readonly SystemWallClockTimeProvider Instance = new();
+
+        public long WallClockTime => DateTime.Now.GetMilliseconds();
+
+        /// <summary>
+        /// No-op for system time provider. Time advances naturally.
+        /// </summary>
+        public void Advance(TimeSpan advance) { }
+    }
+}

core/Mock/TaskSynchronousTopologyDriver.cs

@@ -27,17 +27,18 @@ internal sealed class TaskSynchronousTopologyDriver : IBehaviorTopologyTestDrive
         private readonly IDictionary<TaskId, StreamTask> tasks = new Dictionary<TaskId, StreamTask>();
         private readonly GlobalStateUpdateTask globalTask;
         private readonly IDictionary<string, ExternalProcessorTopologyExecutor> externalProcessorTopologies =
-            new Dictionary<string, ExternalProcessorTopologyExecutor>();        
+            new Dictionary<string, ExternalProcessorTopologyExecutor>();
         private readonly GlobalProcessorContext globalProcessorContext;
         private readonly StreamMetricsRegistry metricsRegistry;
-        
+
         private readonly IDictionary<TaskId, IList<TopicPartition>> partitionsByTaskId =
             new Dictionary<TaskId, IList<TopicPartition>>();
 
         private readonly StreamsProducer producer = null;
         private readonly bool hasGlobalTopology = false;
         private ITopicManager internalTopicManager;
         private readonly IConsumer<byte[],byte[]> repartitionConsumerForwarder;
+        private readonly MockWallClockTimeProvider wallClockTimeProvider;
 
         public bool IsRunning { get; private set; }
 
@@ -113,6 +114,9 @@ public TaskSynchronousTopologyDriver(string clientId, InternalTopologyBuilder to
 
             repartitionConsumerForwarder = this.supplier.GetConsumer(topicConfiguration.ToConsumerConfig("consumer-repartition-forwarder"),
                 null);
+
+            // Initialize mock wall clock time provider for controlling time in tests
+            wallClockTimeProvider = new MockWallClockTimeProvider();
         }
 
         internal StreamTask GetTask(string topicName)
@@ -125,6 +129,13 @@ internal StreamTask GetTask(string topicName)
             {
                 if (builder.GetSourceTopics().Contains(topicName))
                 {
+                    // Make sure partitionsByTaskId has this task's partitions
+                    if (!partitionsByTaskId.ContainsKey(id))
+                    {
+                        var part = new TopicPartition(topicName, 0);
+                        partitionsByTaskId.Add(id, new List<TopicPartition> {part});
+                    }
+
                     task = new StreamTask("thread-0",
                         id,
                         partitionsByTaskId[id],
@@ -134,7 +145,8 @@ internal StreamTask GetTask(string topicName)
                         supplier,
                         producer,
                         new MockChangelogRegister(),
-                        metricsRegistry);
+                        metricsRegistry,
+                        wallClockTimeProvider);
                     task.InitializeStateStores();
                     task.InitializeTopology();
                     task.RestorationIfNeeded();
@@ -163,36 +175,68 @@ private void InitializeInternalTopicManager()
         }
 
         private void ForwardRepartitionTopic(string topic)
+        {
+            ForwardTopicRecordsIfAny(topic);
+        }
+
+        /// <summary>
+        /// Forwards any pending records from an internal topic to downstream tasks.
+        /// Returns true if any records were forwarded, false otherwise.
+        /// </summary>
+        private bool ForwardTopicRecordsIfAny(string topic)
         {
             var records = new List<ConsumeResult<byte[], byte[]>>();
+
             repartitionConsumerForwarder.Subscribe(topic);
+
             ConsumeResult<byte[], byte[]> record = null;
             do
             {
                 record = repartitionConsumerForwarder.Consume();
                 if (record != null)
+                {
                     records.Add(record);
+                }
             } while (record != null);
 
-
             if (records.Any())
             {
+                var task = GetTask(topic);
+
                 var pipe = CreateBuilder(topic).Input(topic, configuration);
                 foreach(var r in records)
                     pipe.Pipe(r.Message.Key, r.Message.Value, r.Message.Timestamp.UnixTimestampMs.FromMilliseconds(), r.Message.Headers);
                 pipe.Flush();
                 pipe.Dispose();
 
                 repartitionConsumerForwarder.Commit(records.Last());
+                repartitionConsumerForwarder.Unsubscribe();
+                return true;
             }
 
             repartitionConsumerForwarder.Unsubscribe();
+            return false;
+        }
+
+        /// <summary>
+        /// Gets all topics that are used for internal communication within the topology.
+        /// These are topics that are both written to (sink) and read from (source) within the topology.
+        /// </summary>
+        private HashSet<string> GetAllInternalCommunicationTopics()
+        {
+            var topicGroups = builder.MakeInternalTopicGroups();
+            var allSinks = topicGroups.Values.SelectMany(t => t.SinkTopics).ToHashSet();
+            var allSources = topicGroups.Values.SelectMany(t => t.SourceTopics).ToHashSet();
+
+            // Internal topics are those that are both written and read within the topology
+            allSinks.IntersectWith(allSources);
+            return allSinks;
         }
 
         private SyncPipeBuilder CreateBuilder(string topicName)
         {
             var task = GetTask(topicName);
-            
+
             if (task == null)
             {
                 if(hasGlobalTopology && builder.GetGlobalTopics().Contains(topicName))
@@ -319,6 +363,77 @@ public void TriggerCommit()
             globalTask?.FlushState();
         }
 
+        public void AdvanceWallClockTime(TimeSpan advance)
+        {
+            // Advance the mock wall clock once at the beginning
+            wallClockTimeProvider.Advance(advance);
+
+            // Get all topics used for internal communication within the topology
+            var internalTopics = GetAllInternalCommunicationTopics();
+
+            // Process until no more progress (like Java's completeAllProcessableWork)
+            // We advance time incrementally inside the loop to allow downstream buffers to flush.
+            // This is needed because buffers schedule punctuations relative to current time,
+            // so buffers created during forwarding need additional time to flush.
+            bool madeProgress;
+            int iterations = 0;
+            const int maxIterations = 100; // Prevent infinite loops
+            do
+            {
+                madeProgress = false;
+                iterations++;
+
+                // 1. FIRST forward records from all internal topics to ensure buffers have data
+                // Keep forwarding until no more records found in any topic
+                bool forwardedAny = false;
+                bool forwardedThisPass;
+                do
+                {
+                    forwardedThisPass = false;
+                    foreach (var topic in internalTopics)
+                    {
+                        if (ForwardTopicRecordsIfAny(topic))
+                        {
+                            madeProgress = true;
+                            forwardedThisPass = true;
+                            forwardedAny = true;
+                        }
+                    }
+                } while (forwardedThisPass);
+
+                // 2. THEN trigger PROCESSING_TIME punctuations for all tasks (including newly created ones)
+                bool punctuatedAny = false;
+                foreach (var task in tasks.Values)
+                {
+                    if (task.PunctuateSystemTime())
+                    {
+                        madeProgress = true;
+                        punctuatedAny = true;
+                    }
+                }
+
+                // 3. If we forwarded records but no punctuations fired, we need to advance
+                // time significantly so downstream buffers can flush. Use the full advance
+                // amount since buffer punctuations may be scheduled far apart.
+                // If punctuations fired or no records were forwarded, use a smaller increment.
+                if (madeProgress)
+                {
+                    if (forwardedAny && !punctuatedAny)
+                    {
+                        // Records went to downstream buffers but their punctuations didn't fire yet
+                        // Advance by full amount to ensure downstream punctuations can fire
+                        wallClockTimeProvider.Advance(advance);
+                    }
+                    else
+                    {
+                        // Normal progress - advance by small amount
+                        wallClockTimeProvider.Advance(TimeSpan.FromMilliseconds(advance.TotalMilliseconds / 10));
+                    }
+                }
+
+            } while (madeProgress && iterations < maxIterations);
+        }
+
         #endregion
     }
 }

core/Mock/TopologyTestDriver.cs

@@ -405,5 +405,44 @@ public IReadOnlyWindowStore<K, V> GetWindowStore<K, V>(string name)
         public bool IsError => behavior.IsError;
 
         #endregion
+
+        #region Time Advancement
+
+        /// <summary>
+        /// Advance wall-clock time by the specified duration, triggering any <see cref="Processors.Public.PunctuationType.PROCESSING_TIME"/>
+        /// punctuations that are scheduled to fire during the advanced time period.
+        /// <para>
+        /// This is analogous to <a href="https://kafka.apache.org/32/javadoc/org/apache/kafka/streams/TopologyTestDriver.html#advanceWallClockTime(java.time.Duration)">
+        /// Java Kafka Streams TopologyTestDriver.advanceWallClockTime</a>.
+        /// </para>
+        /// <para>
+        /// This method is only supported when using <see cref="Mode.SYNC_TASK"/> (the default mode).
+        /// Using this method with <see cref="Mode.ASYNC_CLUSTER_IN_MEMORY"/> will throw <see cref="NotSupportedException"/>.
+        /// </para>
+        /// </summary>
+        /// <param name="advance">The amount of time to advance the wall clock</param>
+        /// <exception cref="NotSupportedException">When used with <see cref="Mode.ASYNC_CLUSTER_IN_MEMORY"/></exception>
+        /// <example>
+        /// <code>
+        /// using (var driver = new TopologyTestDriver(topology, config))
+        /// {
+        ///     var inputTopic = driver.CreateInputTopic&lt;string, string&gt;("input");
+        ///     var outputTopic = driver.CreateOutputTopic&lt;string, string&gt;("output");
+        ///
+        ///     inputTopic.PipeInput("key", "value");
+        ///
+        ///     // Advance time to trigger PROCESSING_TIME punctuations
+        ///     driver.AdvanceWallClockTime(TimeSpan.FromMilliseconds(500));
+        ///
+        ///     var result = outputTopic.ReadKeyValue();
+        /// }
+        /// </code>
+        /// </example>
+        public void AdvanceWallClockTime(TimeSpan advance)
+        {
+            behavior.AdvanceWallClockTime(advance);
+        }
+
+        #endregion
     }
 }