mac: Several minor fixes to allocator shim.

base/debug/thread_heap_usage_tracker_unittest.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "base/debug/thread_heap_usage_tracker.h"
 
 #include <map>
 
 #include "base/allocator/allocator_shim.h"
 #include "base/allocator/features.h"
 #include "testing/gtest/include/gtest/gtest.h"
 
+#if defined(OS_MACOSX)
+#include "base/allocator/allocator_interception_mac.h"
+#endif
+
 namespace base {
 namespace debug {
 
 namespace {
 
 class TestingThreadHeapUsageTracker : public ThreadHeapUsageTracker {
  public:
   using ThreadHeapUsageTracker::DisableHeapTrackingForTesting;
   using ThreadHeapUsageTracker::EnsureTLSInitialized;
   using ThreadHeapUsageTracker::GetDispatchForTesting;
@@ skipping 520 matching lines @@
   estimate = MockGetSizeEstimate(alloc);
   ASSERT_TRUE(estimate == 0 || estimate >= kAllocSize);
 
   alloc = MockRealloc(alloc, kAllocSize);
   estimate = MockGetSizeEstimate(alloc);
   ASSERT_TRUE(estimate == 0 || estimate >= kAllocSize);
   MockFree(alloc);
 }
 
 #if BUILDFLAG(USE_EXPERIMENTAL_ALLOCATOR_SHIM)
-TEST(ThreadHeapUsageShimTest, HooksIntoMallocWhenShimAvailable) {
+class ThreadHeapUsageShimTest : public testing::Test {
 #if defined(OS_MACOSX)
-  allocator::InitializeAllocatorShim();
+  void SetUp() override { allocator::InitializeAllocatorShim(); }
+  void TearDown() override { allocator::UninterceptMallocZonesForTesting(); }
 #endif
+};
 
+TEST_F(ThreadHeapUsageShimTest, HooksIntoMallocWhenShimAvailable) {
   ASSERT_FALSE(ThreadHeapUsageTracker::IsHeapTrackingEnabled());
 
   ThreadHeapUsageTracker::EnableHeapTracking();
 
   ASSERT_TRUE(ThreadHeapUsageTracker::IsHeapTrackingEnabled());
 
   const size_t kAllocSize = 9993;
   // This test verifies that the scoped heap data is affected by malloc &
   // free only when the shim is available.
   ThreadHeapUsageTracker usage_tracker;
@@ skipping 23 matching lines @@
   EXPECT_LE(u2.free_ops + 1, u3.free_ops);
 
   TestingThreadHeapUsageTracker::DisableHeapTrackingForTesting();
 
   ASSERT_FALSE(ThreadHeapUsageTracker::IsHeapTrackingEnabled());
 }
 #endif  // BUILDFLAG(USE_EXPERIMENTAL_ALLOCATOR_SHIM)
 
 }  // namespace debug
 }  // namespace base