3 files changed, 8 insertions, 18 deletions
diff --git a/xlators/cluster/dht/src/Makefile.am b/xlators/cluster/dht/src/Makefile.am
index ab58affe2b4..8d02749f4d9 100644
--- a/xlators/cluster/dht/src/Makefile.am
+++ b/xlators/cluster/dht/src/Makefile.am
@@ -36,7 +36,7 @@ uninstall-local:
 install-data-hook:
 	ln -sf dht.so $(DESTDIR)$(xlatordir)/distribute.so
 
-#### UNIT TESTS #####
+if UNITTEST
 CLEANFILES += *.gcda *.gcno *_xunit.xml
 noinst_PROGRAMS =
 TESTS =
@@ -49,3 +49,4 @@ dht_layout_unittest_CFLAGS = $(AM_CFLAGS) $(UNITTEST_CFLAGS)
 dht_layout_unittest_LDFLAGS = $(UNITTEST_LDFLAGS)
 noinst_PROGRAMS += dht_layout_unittest
 TESTS += dht_layout_unittest
+endif
diff --git a/xlators/cluster/dht/src/dht-layout.c b/xlators/cluster/dht/src/dht-layout.c
index fa29b6bfa28..757ec731d26 100644
--- a/xlators/cluster/dht/src/dht-layout.c
+++ b/xlators/cluster/dht/src/dht-layout.c
@@ -19,6 +19,7 @@
 #include "dht-common.h"
 #include "byte-order.h"
 #include "dht-messages.h"
+#include "unittest/unittest.h"
 
 
 #define layout_base_size (sizeof (dht_layout_t))
@@ -27,18 +28,6 @@
 
 #define layout_size(cnt) (layout_base_size + (cnt * layout_entry_size))
 
-#include <cmockery/pbc.h>
-#include <cmockery/cmockery_override.h>
-
-// Change GF_CALLOC and GF_FREE to use
-// cmockery2 memory allocation versions
-#ifdef UNIT_TESTING
-#undef GF_CALLOC
-#define GF_CALLOC(n, s, t) test_calloc(n, s)
-#undef GF_FREE
-#define GF_FREE test_free
-#endif
-
 dht_layout_t *
 dht_layout_new (xlator_t *this, int cnt)
 {
diff --git a/xlators/cluster/dht/src/unittest/dht_layout_unittest.c b/xlators/cluster/dht/src/unittest/dht_layout_unittest.c
index b5233d235d0..3702af366f9 100644
--- a/xlators/cluster/dht/src/unittest/dht_layout_unittest.c
+++ b/xlators/cluster/dht/src/unittest/dht_layout_unittest.c
@@ -12,12 +12,12 @@
 #include "logging.h"
 #include "xlator.h"
 
+#include <inttypes.h>
 #include <stdarg.h>
 #include <stddef.h>
 #include <setjmp.h>
-#include <inttypes.h>
-#include <cmockery/pbc.h>
-#include <cmockery/cmockery.h>
+#include <cmocka_pbc.h>
+#include <cmocka.h>
 
 /*
  * Helper functions
@@ -116,9 +116,9 @@ test_dht_layout_new(void **state)
 }
 
 int main(void) {
-    const UnitTest tests[] = {
+    const struct CMUnitTest xlator_dht_layout_tests[] = {
         unit_test(test_dht_layout_new),
     };
 
-    return run_tests(tests, "xlator_dht_layout");
+    return cmocka_run_group_tests(xlator_dht_layout_tests, NULL, NULL);
 }