From bf73a9a43daca9cb74f7f968c83a5efbfe13e3a5 Mon Sep 17 00:00:00 2001
From: Hyeonseo Yang <jourhyang123@gmail.com>
Date: Thu, 2 Jan 2025 22:14:06 +0900
Subject: [PATCH] :bug: Fix cuda link errors in test

---
 CMakeLists.txt           | 32 ++++++++++++++++++++------------
 tests/cpp/CMakeLists.txt | 27 +++++++++++++++++++++++++--
 tests/cpp/test_ipc.cpp   |  8 +++++---
 3 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 1ffc4226..5647aebd 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -82,31 +82,39 @@ install(TARGETS craftground_native LIBRARY DESTINATION craftground)
 option(BUILD_TESTS "Build tests" ON)
 if(BUILD_TESTS)
     enable_testing()
-    pybind11_add_module(craftground STATIC ${CRAFTGROUND_PY_SOURCES})
+    add_library(craftground STATIC ${CRAFTGROUND_PY_SOURCES})
+    find_package(Python3 REQUIRED COMPONENTS Development)
+    target_include_directories(
+        craftground PRIVATE
+        ${CMAKE_CURRENT_SOURCE_DIR}/pybind11/include
+        ${Python3_INCLUDE_DIRS}
+    )
     if(APPLE)
         if(Torch_FOUND)
-            target_link_libraries(craftground PRIVATE "${TORCH_LIBRARIES}")
+            target_link_libraries(craftground PUBLIC "${TORCH_LIBRARIES}")
             set(LIBTORCH_DIR ${CMAKE_SOURCE_DIR}/src/cpp/libtorch)
-            target_link_libraries(craftground PRIVATE
+            target_link_libraries(craftground PUBLIC
                 ${LIBTORCH_DIR}/libtorch_cpu_minimal.a
                 ${LIBTORCH_DIR}/libc10.a
             )
-            target_include_directories(craftground PRIVATE "${TORCH_INCLUDE_DIRS}")
-            target_compile_definitions(craftground PRIVATE HAS_TORCH)
-            target_compile_options(craftground PRIVATE -ffunction-sections -fdata-sections)
-            target_link_options(craftground PRIVATE -Wl,-dead_strip)
+            target_include_directories(craftground PUBLIC "${TORCH_INCLUDE_DIRS}")
+            target_compile_definitions(craftground PUBLIC HAS_TORCH)
+            target_compile_options(craftground PUBLIC -ffunction-sections -fdata-sections)
+            # target_link_options(craftground PUBLIC -Wl,-dead_strip)
         endif()
     elseif(CUDAToolkit_FOUND)
-        target_include_directories(craftground PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
-        target_link_libraries(craftground PRIVATE ${CUDAToolkit_LIBRARIES})
+        message(STATUS "CUDA is available in tests")
+        target_include_directories(craftground PUBLIC ${CUDAToolkit_INCLUDE_DIRS})
+        target_link_libraries(craftground PUBLIC CUDA::cudart CUDA::cudart_static)
     endif()
-    target_include_directories(craftground PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src/cpp)
+    target_link_options(craftground PRIVATE "-Wl,--whole-archive" "-Wl,--no-whole-archive")
+    target_include_directories(craftground PUBLIC ${CMAKE_CURRENT_SOURCE_DIR}/src/cpp)
 
-    target_compile_options(craftground PRIVATE ${CRAFTGROUND_PY_COMPILE_OPTIONS})
+    target_compile_options(craftground PUBLIC ${CRAFTGROUND_PY_COMPILE_OPTIONS})
 
     target_compile_definitions(
         craftground
-        PRIVATE VERSION_INFO=${PRIVATE_VERSION_INFO}
+        PUBLIC VERSION_INFO=${PRIVATE_VERSION_INFO}
     )
     add_subdirectory(tests)
 endif()
diff --git a/tests/cpp/CMakeLists.txt b/tests/cpp/CMakeLists.txt
index 8566931f..65b71f4e 100644
--- a/tests/cpp/CMakeLists.txt
+++ b/tests/cpp/CMakeLists.txt
@@ -2,10 +2,33 @@
 find_package(GTest REQUIRED)
 include_directories(${GTEST_INCLUDE_DIRS})
 
+find_package(CUDAToolkit REQUIRED)
+
+find_package(Python3 REQUIRED COMPONENTS Development)
+
 # Build Test directory
 add_executable(test_ipc test_ipc.cpp)
-target_include_directories(test_ipc PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/../../src/cpp)
-target_link_libraries(test_ipc PRIVATE craftground GTest::GTest GTest::Main)
+
+target_include_directories(
+    test_ipc PRIVATE
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../src/cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../pybind11/include
+    ${Python3_INCLUDE_DIRS}
+)
+target_include_directories(test_ipc PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
+message(STATUS "CUDA is available in tests")
+message(STATUS "CUDa toolkit libraries: ${CUDAToolkit_LIBRARY_DIR}")
+target_link_libraries(
+    test_ipc PRIVATE
+    GTest::GTest
+    GTest::Main
+    Python3::Python
+    CUDA::cudart
+    CUDA::cudart_static
+    craftground
+)
+target_compile_options(test_ipc PRIVATE -D_GLIBCXX_USE_CXX11_ABI=1)
+
 
 # Integrate with CTest
 add_test(NAME CppTests COMMAND test_ipc)
diff --git a/tests/cpp/test_ipc.cpp b/tests/cpp/test_ipc.cpp
index 48112c58..549c009b 100644
--- a/tests/cpp/test_ipc.cpp
+++ b/tests/cpp/test_ipc.cpp
@@ -1,11 +1,13 @@
 #include <gtest/gtest.h>
+#include <pybind11/pybind11.h>
 #include "ipc.h"
 
 TEST(IPCModuleTest, AddFunction) {
-    EXPECT_EQ(add(2, 3), 5);
+    EXPECT_EQ(initialize_from_mach_port(2, 3, 4), py::none());
 }
 
 TEST(IPCModuleTest, ExceptionHandling) {
-    EXPECT_THROW(safe_divide(1, 0), std::invalid_argument);
+    EXPECT_THROW(
+        mtl_tensor_from_cuda_mem_handle("", 0, 2), std::invalid_argument
+    );
 }
-