test(precompute): add a wall-clock override hook to PrecomputeEngine (#537)

Stacked on the #474 fix. Worker already has a fake-clock override
(set_now_ms_fn) for its own unit tests, but it's #[cfg(test)]-gated and
invisible to integration tests under tests/, which link the library
normally. Widen it to pub(crate) and add PrecomputeEngine::with_now_ms_fn,
a builder that threads a shared clock override down to every spawned
worker, so an integration test can drive the wall-clock fallback
deterministically through the real ingest pipeline instead of racing a
real clock with sleeps.

No behavioral change: the override defaults to None and every worker keeps
its default SystemTime::now-backed clock unless a caller opts in.
Production code never calls with_now_ms_fn.

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>

Author: milindsrivastava1997

asap-query-engine/src/precompute_engine/engine.rs

@@ -74,6 +74,10 @@ pub struct PrecomputeEngine {
     receivers: Option<Vec<mpsc::Receiver<WorkerMessage>>>,
     /// Shared ingest agg_configs, swappable at runtime.
     ingest_agg_configs: Arc<ArcSwap<Vec<Arc<AggregationConfig>>>>,
+    /// Test-support wall-clock override, applied to every spawned worker.
+    /// See `with_now_ms_fn`. `None` in production — each worker keeps its
+    /// default `SystemTime::now`-backed clock.
+    now_ms_fn: Option<Arc<dyn Fn() -> i64 + Send + Sync>>,
 }
 
 impl PrecomputeEngine {
@@ -119,6 +123,7 @@ impl PrecomputeEngine {
             senders,
             receivers: Some(receivers),
             ingest_agg_configs,
+            now_ms_fn: None,
         }
     }
 
@@ -127,6 +132,16 @@ impl PrecomputeEngine {
         self.diagnostics.clone()
     }
 
+    /// Test-support builder: override the wall-clock source used by every
+    /// worker's wall-clock fallback (`flush_all`). Lets integration tests
+    /// drive the fallback deterministically through the real ingest pipeline
+    /// instead of racing a real clock with `tokio::time::sleep`. Must be
+    /// called before `run()`. Production code never calls this.
+    pub fn with_now_ms_fn(mut self, f: impl Fn() -> i64 + Send + Sync + 'static) -> Self {
+        self.now_ms_fn = Some(Arc::new(f));
+        self
+    }
+
     /// Return a handle for applying runtime config updates to this engine.
     /// Must be called before `run()`.
     pub fn handle(&self) -> PrecomputeEngineHandle {
@@ -160,7 +175,7 @@ impl PrecomputeEngine {
         // Spawn workers
         let mut worker_handles = Vec::with_capacity(num_workers);
         for (id, rx) in receivers.into_iter().enumerate() {
-            let worker = Worker::new(
+            let mut worker = Worker::new(
                 id,
                 rx,
                 self.output_sink.clone(),
@@ -177,6 +192,10 @@ impl PrecomputeEngine {
                 self.diagnostics.worker_watermarks[id].clone(),
                 self.diagnostics.worker_watermarks.to_vec(),
             );
+            if let Some(now_ms_fn) = &self.now_ms_fn {
+                let now_ms_fn = now_ms_fn.clone();
+                worker.set_now_ms_fn(Box::new(move || now_ms_fn()));
+            }
             let handle = tokio::spawn(async move {
                 worker.run().await;
             });

asap-query-engine/src/precompute_engine/worker.rs

@@ -112,7 +112,8 @@ pub struct Worker {
     wall_clock_grace_period_ms: i64,
     /// Injectable clock returning current wall-clock time in milliseconds
     /// since the unix epoch. Production uses `SystemTime::now`; tests
-    /// override with a deterministic fake via `set_now_ms_fn`.
+    /// override with a deterministic fake via `set_now_ms_fn` (directly, or
+    /// via `PrecomputeEngine::with_now_ms_fn` for integration tests).
     now_ms_fn: Box<dyn Fn() -> i64 + Send + Sync>,
 }
 
@@ -154,12 +155,14 @@ impl Worker {
         }
     }
 
-    /// Test/diagnostic-only setter for the wall-clock source. Replaces the
-    /// default `SystemTime::now`-backed clock with a deterministic fake so
-    /// unit tests can drive the wall-clock fallback in `flush_all` without
-    /// `std::thread::sleep`. Production code never calls this.
-    #[cfg(test)]
-    pub fn set_now_ms_fn(&mut self, f: Box<dyn Fn() -> i64 + Send + Sync>) {
+    /// Test-support setter for the wall-clock source. Replaces the default
+    /// `SystemTime::now`-backed clock with a deterministic fake so tests can
+    /// drive the wall-clock fallback in `flush_all` without real sleeping.
+    /// Crate-visible (not `#[cfg(test)]`-gated) so `PrecomputeEngine::
+    /// with_now_ms_fn` can reach it from integration tests in `tests/`,
+    /// which link the library normally and don't see `#[cfg(test)]` items.
+    /// Production code never calls this.
+    pub(crate) fn set_now_ms_fn(&mut self, f: Box<dyn Fn() -> i64 + Send + Sync>) {
         self.now_ms_fn = f;
     }