diff --git a/be/test/agent/resource_group_usage_recorder_test.cpp b/be/test/agent/resource_group_usage_recorder_test.cpp index ca0c61644907d9..8084f29a949361 100644 --- a/be/test/agent/resource_group_usage_recorder_test.cpp +++ b/be/test/agent/resource_group_usage_recorder_test.cpp @@ -24,10 +24,14 @@ TEST(ResourceGroupUsageRecorderTest, test_get_resource_group_usages) { const size_t num_cores = CpuInfo::num_cores(); auto& exec_env = *ExecEnv::GetInstance(); + // Save original workgroup_manager to restore at end (otherwise subsequent tests fail) + auto original_wg_manager = std::move(exec_env._workgroup_manager); + workgroup::PipelineExecutorSetConfig executors_manager_opts( CpuInfo::num_cores(), num_cores, num_cores, num_cores, CpuInfo::get_core_ids(), true, config::enable_resource_group_cpu_borrowing, StarRocksMetrics::instance()->get_pipeline_executor_metrics()); exec_env._workgroup_manager = std::make_unique(std::move(executors_manager_opts)); + ASSERT_TRUE(exec_env._workgroup_manager->start().ok()); workgroup::DefaultWorkGroupInitialization default_workgroup_init; auto default_wg = exec_env.workgroup_manager()->get_default_workgroup(); @@ -43,6 +47,9 @@ TEST(ResourceGroupUsageRecorderTest, test_get_resource_group_usages) { ASSERT_EQ(group_usages[0].mem_pool, workgroup::WorkGroup::DEFAULT_MEM_POOL); ASSERT_EQ(group_usages[0].mem_limit_bytes, default_wg->mem_limit_bytes()); ASSERT_EQ(group_usages[0].mem_pool_mem_limit_bytes, default_wg->mem_limit_bytes()); + + // Restore original workgroup_manager + exec_env._workgroup_manager = std::move(original_wg_manager); } } // namespace starrocks \ No newline at end of file diff --git a/be/test/exec/pipeline/pipeline_file_scan_node_test.cpp b/be/test/exec/pipeline/pipeline_file_scan_node_test.cpp index 027055d319ba38..43860f29e47765 100644 --- a/be/test/exec/pipeline/pipeline_file_scan_node_test.cpp +++ b/be/test/exec/pipeline/pipeline_file_scan_node_test.cpp @@ -81,6 +81,7 @@ class PipeLineFileScanNodeTest : public ::testing::Test { _fragment_ctx->set_runtime_state( std::make_unique(_request.params.query_id, _request.params.fragment_instance_id, _request.query_options, _request.query_globals, _exec_env)); + _fragment_ctx->set_workgroup(ExecEnv::GetInstance()->workgroup_manager()->get_default_workgroup()); _fragment_future = _fragment_ctx->finish_future(); _runtime_state = _fragment_ctx->runtime_state();