Add ASSERT_REF to check for reference counts of af::arrays

umar456 · Aug 17, 2021 · 95d11c3 · 95d11c3
1 parent 75bc1d5
commit 95d11c3
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/test/arrayfire_test.cpp b/test/arrayfire_test.cpp
@@ -1634,6 +1634,23 @@ ::testing::AssertionResult assertArrayNear(
                            bbb, maxAbsDiff);
 }
 
+::testing::AssertionResult assertRefEq(std::string hA_name,
+                                       std::string expected_name,
+                                       const af::array &a, int expected) {
+    int count = 0;
+    af_get_data_ref_count(&count, a.get());
+    if (count != expected) {
+        std::stringstream ss;
+        ss << "Incorrect reference count:\nExpected: " << expected << "\n"
+           << std::setw(8) << hA_name << ": " << count;
+
+        return ::testing::AssertionFailure() << ss.str();
+
+    } else {
+        return ::testing::AssertionSuccess();
+    }
+}
+
 #define INSTANTIATE(To)                                                        \
     template std::string printContext(                                         \
         const std::vector<To> &hGold, std::string goldName,                    \

diff --git a/test/testHelpers.hpp b/test/testHelpers.hpp
@@ -360,6 +360,10 @@ ::testing::AssertionResult assertArrayNear(
     std::string maxAbsDiffName, const std::vector<T> &hA, af::dim4 aDims,
     const af_array b, float maxAbsDiff);
 
+::testing::AssertionResult assertRefEq(std::string hA_name,
+                                       std::string expected_name,
+                                       const af::array &a, int expected);
+
 /// Checks if the C-API arrayfire function returns successfully
 ///
 /// \param[in] CALL This is the arrayfire C function
@@ -430,6 +434,9 @@ ::testing::AssertionResult assertArrayNear(
     ASSERT_PRED_FORMAT4(assertArrayNear, EXPECTED_VEC, EXPECTED_ARR_DIMS,  \
                         ACTUAL_ARR, MAX_ABSDIFF)
 
+#define ASSERT_REF(arr, expected) \
+    ASSERT_PRED_FORMAT2(assertRefEq, arr, expected)
+
 #if defined(USE_MTX)
 ::testing::AssertionResult mtxReadSparseMatrix(af::array &out,
                                                const char *fileName);