From 8115dd708ae80d57ccfe14df5f392d87e029cd74 Mon Sep 17 00:00:00 2001
From: discord9 <discord9@163.com>
Date: Sun, 25 Oct 2026 03:52:17 +0800
Subject: [PATCH] test(promql): restore the shared pool fixture for metric
 lifecycle tests

Signed-off-by: discord9 <discord9@163.com>
---
 .../src/extension_plan/fused_series/tests.rs       | 14 ++++++++++++--
 1 file changed, 12 insertions(+), 2 deletions(-)

diff --git a/src/promql/src/extension_plan/fused_series/tests.rs b/src/promql/src/extension_plan/fused_series/tests.rs
index 59e21a83d66..4a801aa8985 100644
--- a/src/promql/src/extension_plan/fused_series/tests.rs
+++ b/src/promql/src/extension_plan/fused_series/tests.rs
@@ -47,7 +47,8 @@ use datafusion::datasource::memory::MemorySourceConfig;
 use datafusion::datasource::source::DataSourceExec;
 use datafusion::error::{DataFusionError, Result as DataFusionResult};
 use datafusion::execution::context::TaskContext;
-use datafusion::execution::memory_pool::UnboundedMemoryPool;
+use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryPool, UnboundedMemoryPool};
+use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 use datafusion::logical_expr::{EmptyRelation, LogicalPlan, UserDefinedLogicalNodeCore};
 use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr};
 use datafusion::physical_optimizer::PhysicalOptimizerRule;
@@ -55,7 +56,7 @@ use datafusion::physical_optimizer::ensure_requirements::EnsureRequirements;
 use datafusion::physical_plan::execution_plan::{
     Boundedness, ChildrenPropertiesMode, EmissionType, ReplaceChildrenOptions,
 };
-use datafusion::physical_plan::metrics::{ExecutionPlanMetricsSet, MetricValue, MetricsSet};
+use datafusion::physical_plan::metrics::{MetricValue, MetricsSet};
 use datafusion::physical_plan::{
     DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, PlanProperties,
     RecordBatchStream, SendableRecordBatchStream, collect, displayable,
@@ -375,6 +376,15 @@ fn memory_exec_partitioned(partitions: Vec<Vec<RecordBatch>>) -> Arc<dyn Executi
     memory_exec_partitioned_with(frame_schema(), partitions)
 }
 
+/// A task context whose memory pool is `pool`, so the execution's reservations can be observed.
+fn task_context(pool: &Arc<GreedyMemoryPool>) -> Arc<TaskContext> {
+    let runtime = RuntimeEnvBuilder::new()
+        .with_memory_pool(Arc::clone(pool) as Arc<dyn MemoryPool>)
+        .build()
+        .expect("the runtime environment builds");
+    Arc::new(TaskContext::default().with_runtime(Arc::new(runtime)))
+}
+
 /// A physical input whose one partition's stream `make` builds.
 fn staged(
     schema: SchemaRef,