rbtree: use failtest to check handling of allocation failures.

[ccan] / ccan / rbtree / test / run.c

diff --git a/ccan/rbtree/test/run.c b/ccan/rbtree/test/run.c
index 11d0ed78b015f03b252986ab4ea4cead5cde7ae9..e1ce38730cd0e1f9c3c41bc3f25a4ac1cf9481b2 100644 (file)
--- a/ccan/rbtree/test/run.c
+++ b/ccan/rbtree/test/run.c
@@ -1,8 +1,25 @@
+#include <ccan/failtest/failtest.h>
 #include <ccan/rbtree/rbtree.c>
 #include <ccan/tap/tap.h>
 #include <ccan/talloc/talloc.h>
 #include <string.h>
 
+/* We want to test talloc failure paths. */
+static void *my_malloc(size_t size)
+{
+       return malloc(size);
+}
+
+static void my_free(void *ptr)
+{
+       free(ptr);
+}
+
+static void *my_realloc(void *ptr, size_t size)
+{
+       return realloc(ptr, size);
+}
+
 static void *insert_callback(void *param, void *data)
 {
        ok1(data == param);
@@ -12,11 +29,13 @@ static void *insert_callback(void *param, void *data)
 int main(void)
 {
        trbt_tree_t *rb;
-       void *ctx = talloc_init("toplevel");
+       void *ctx = talloc_strdup(NULL, "toplevel");
        char *data, *data2;
 
        /* This is how many tests you plan to run */
-       plan_tests(18);
+       plan_tests(19);
+
+       talloc_set_allocator(my_malloc, my_free, my_realloc);
 
        rb = trbt_create(ctx, 0);
        ok1(rb);
@@ -55,16 +74,23 @@ int main(void)
        /* Insert with callback on existing. */
        trbt_insert32_callback(rb, 0, insert_callback, data2);
        ok1(strcmp(trbt_lookup32(rb, 0), "insert_callback") == 0);
+       talloc_free(data2);
 
        /* Delete. */
+       data2 = trbt_lookup32(rb, 1);
        trbt_delete32(rb, 1);
        ok1(trbt_lookup32(rb, 1) == NULL);
        ok1(trbt_lookup32(rb, 0));
+       talloc_free(data2);
 
        /* This should free everything. */
        talloc_free(trbt_lookup32(rb, 0));
        talloc_free(rb);
 
+       /* No memory leaks? */
+       ok1(talloc_total_blocks(ctx) == 1);
+       talloc_free(ctx);
+
        /* This exits depending on whether all tests passed */
-       return exit_status();
+       failtest_exit(exit_status());
 }