From 988771437f253193637bb631d83ce689286a8acf Mon Sep 17 00:00:00 2001
From: Michael Schellenberger Costa <miscco@nvidia.com>
Date: Fri, 1 Mar 2024 14:53:11 +0100
Subject: [PATCH] Replace usages of raw `get_upstream` with
 `get_upstream_resource()`

We want to get rid of raw memory resources so move to the new interface instead
---
 cpp/test/core/handle.cpp | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

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 <rmm/device_buffer.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <cuda_runtime.h>
 
@@ -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<rmm::mr::device_memory_resource>{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();