diff --git a/cpp/test/core/handle.cpp b/cpp/test/core/handle.cpp index 0b0b4b54ab..1c3bb8d846 100644 --- a/cpp/test/core/handle.cpp +++ b/cpp/test/core/handle.cpp @@ -25,6 +25,7 @@ #include #include #include +#include #include @@ -281,7 +282,8 @@ TEST(Raft, WorkspaceResource) raft::handle_t handle; // The returned resource is always a limiting adaptor - auto* orig_mr = resource::get_workspace_resource(handle)->get_upstream(); + rmm::device_async_resource_ref orig_mr{ + resource::get_workspace_resource(handle)->get_upstream_resource()}; // Let's create a pooled resource auto pool_mr = std::shared_ptr{new rmm::mr::pool_memory_resource( @@ -292,11 +294,11 @@ TEST(Raft, WorkspaceResource) // Replace the resource resource::set_workspace_resource(handle, pool_mr, max_size); - auto new_mr = resource::get_workspace_resource(handle); + rmm::device_async_resource_ref new_mr{resource::get_workspace_resource(handle)}; // By this point, the orig_mr likely points to a non-existent resource; don't dereference! ASSERT_NE(orig_mr, new_mr); - ASSERT_EQ(pool_mr.get(), new_mr->get_upstream()); + ASSERT_EQ(rmm::device_async_resource_ref{pool_mr.get()}, new_mr); // We can safely reset pool_mr, because the shared_ptr to the pool memory stays in the resource pool_mr.reset();