From 184128fd9e5bfcf9b9a1c1021cf1d8db03147c28 Mon Sep 17 00:00:00 2001
From: Benjamin Feng <benjamin.feng@glassdoor.com>
Date: Thu, 30 Jan 2020 00:14:17 -0600
Subject: [PATCH] Fix testing.allocator wiring

---
 lib/std/special/test_runner.zig          | 10 ++++------
 lib/std/testing/leak_count_allocator.zig |  3 ++-
 2 files changed, 6 insertions(+), 7 deletions(-)

diff --git a/lib/std/special/test_runner.zig b/lib/std/special/test_runner.zig
index 9a2d68a73d..07c46ee43c 100644
--- a/lib/std/special/test_runner.zig
+++ b/lib/std/special/test_runner.zig
@@ -14,12 +14,6 @@ pub fn main() anyerror!void {
 
     for (test_fn_list) |test_fn, i| {
         std.testing.base_allocator_instance.reset();
-        defer {
-            std.testing.allocator_instance.validate() catch |err| switch (err) {
-                error.Leak => std.debug.panic("", .{}),
-                else => std.debug.panic("error.{}", .{@errorName(err)}),
-            };
-        }
 
         var test_node = root_node.start(test_fn.name, null);
         test_node.activate();
@@ -30,6 +24,10 @@ pub fn main() anyerror!void {
         if (test_fn.func()) |_| {
             ok_count += 1;
             test_node.end();
+            std.testing.allocator_instance.validate() catch |err| switch (err) {
+                error.Leak => std.debug.panic("", .{}),
+                else => std.debug.panic("error.{}", .{@errorName(err)}),
+            };
             if (progress.terminal == null) std.debug.warn("OK\n", .{});
         } else |err| switch (err) {
             error.SkipZigTest => {
diff --git a/lib/std/testing/leak_count_allocator.zig b/lib/std/testing/leak_count_allocator.zig
index 6cdcfa2678..1d65b8c560 100644
--- a/lib/std/testing/leak_count_allocator.zig
+++ b/lib/std/testing/leak_count_allocator.zig
@@ -23,10 +23,11 @@ pub const LeakCountAllocator = struct {
 
     fn realloc(allocator: *std.mem.Allocator, old_mem: []u8, old_align: u29, new_size: usize, new_align: u29) ![]u8 {
         const self = @fieldParentPtr(LeakCountAllocator, "allocator", allocator);
+        var data = try self.internal_allocator.reallocFn(self.internal_allocator, old_mem, old_align, new_size, new_align);
         if (old_mem.len == 0) {
             self.count += 1;
         }
-        return self.internal_allocator.reallocFn(self.internal_allocator, old_mem, old_align, new_size, new_align);
+        return data;
     }
 
     fn shrink(allocator: *std.mem.Allocator, old_mem: []u8, old_align: u29, new_size: usize, new_align: u29) []u8 {