From 2db3dfc803e1f081635895b53984987ee47ff7a1 Mon Sep 17 00:00:00 2001
From: Eric Cousineau <eric.cousineau@tri.global>
Date: Tue, 14 Apr 2020 14:53:53 -0400
Subject: [PATCH] test_smart_ptr: Add test for reclaim_from_cpp logic

---
 tests/test_smart_ptr.cpp | 23 +++++++++++++++++++++++
 tests/test_smart_ptr.py  | 14 ++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/tests/test_smart_ptr.cpp b/tests/test_smart_ptr.cpp
index f7c46339c0..b4e2c195c7 100644
--- a/tests/test_smart_ptr.cpp
+++ b/tests/test_smart_ptr.cpp
@@ -462,4 +462,27 @@ TEST_SUBMODULE(smart_ptr, m) {
         [](std::shared_ptr<SharedPtrHeld> obj) {
             return obj != nullptr && obj->value == 10;
         });
+
+    // Test passing ownership of registered, but unowned, C++ instances back to
+    // Python. This happens when a raw pointer is passed first, and then
+    // ownership is transfered.
+    struct UniquePtrHeldContainer {
+        UniquePtrHeldContainer() {
+            value_.reset(new UniquePtrHeld(10));
+        }
+        UniquePtrHeld* get() const {
+            return value_.get();
+        }
+        using Ptr = std::unique_ptr<UniquePtrHeld>;
+        Ptr reset(Ptr to) {
+            Ptr from = std::move(value_);
+            value_ = std::move(to);
+            return from;
+        }
+        std::unique_ptr<UniquePtrHeld> value_;
+    };
+    py::class_<UniquePtrHeldContainer>(m, "UniquePtrHeldContainer")
+        .def(py::init())
+        .def("get", &UniquePtrHeldContainer::get, py::return_value_policy::reference_internal)
+        .def("reset", &UniquePtrHeldContainer::reset);
 }
diff --git a/tests/test_smart_ptr.py b/tests/test_smart_ptr.py
index a79a84a9ac..63d9171b59 100644
--- a/tests/test_smart_ptr.py
+++ b/tests/test_smart_ptr.py
@@ -335,3 +335,17 @@ def test_unique_ptr_overload_fail():
     out = m.unique_ptr_overload(obj, m.SecondT())
     assert out["obj"] is obj
     assert out["overload"] == 2
+
+
+def test_unique_ptr_held_container_from_cpp():
+
+    def check_reset(obj_new):
+        c = m.UniquePtrHeldContainer()
+        obj = c.get()
+        assert obj is not None
+        obj_ref = c.reset(obj_new)
+        assert obj is obj_ref
+        assert c.get() is obj_new
+
+    check_reset(obj_new=m.UniquePtrHeld(10))
+    check_reset(obj_new=None)