From 9168fbdec87ddcf5c6ab1ac67904f2f170d84899 Mon Sep 17 00:00:00 2001
From: Bryan Bernhart <bryan.bernhart@intel.com>
Date: Fri, 16 Sep 2022 14:49:58 -0700
Subject: [PATCH] Refactor  harness to keep per test allocator descriptions
 separate.

---
 src/tests/D3D12Test.cpp                       | 11 +--------
 src/tests/D3D12Test.h                         |  2 +-
 .../D3D12EventTraceReplay.cpp                 |  7 ++++--
 .../end2end/D3D12ResourceAllocatorTests.cpp   | 23 +++++++++++++++++--
 4 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/src/tests/D3D12Test.cpp b/src/tests/D3D12Test.cpp
index 0fdc5faa1..5b915c1fd 100644
--- a/src/tests/D3D12Test.cpp
+++ b/src/tests/D3D12Test.cpp
@@ -109,7 +109,7 @@ namespace gpgmm::d3d12 {
         GPGMMTestBase::TearDown();
     }
 
-    ALLOCATOR_DESC D3D12TestBase::CreateBasicAllocatorDesc(bool isPrefetchAllowed) const {
+    ALLOCATOR_DESC D3D12TestBase::CreateBasicAllocatorDesc() const {
         ALLOCATOR_DESC desc = {};
 
         // Required parameters.
@@ -117,15 +117,6 @@ namespace gpgmm::d3d12 {
         desc.Device = mDevice;
         desc.ResourceHeapTier = mResourceHeapTier;
 
-        // Pre-fetching is enabled by default. However for testing purposes, pre-fetching changes
-        // expectations that check GPU memory usage and needs to be tested in isolation.
-        if (!isPrefetchAllowed) {
-            desc.Flags |= ALLOCATOR_FLAG_DISABLE_MEMORY_PREFETCH;
-        }
-
-        // Make sure leak detection is always enabled.
-        desc.Flags |= gpgmm::d3d12::ALLOCATOR_FLAG_NEVER_LEAK_MEMORY;
-
         desc.MinLogLevel = GetMessageSeverity(GetLogLevel());
 
         if (IsDumpAllEventsEnabled()) {
diff --git a/src/tests/D3D12Test.h b/src/tests/D3D12Test.h
index 0e35076e4..37938576c 100644
--- a/src/tests/D3D12Test.h
+++ b/src/tests/D3D12Test.h
@@ -40,7 +40,7 @@ namespace gpgmm::d3d12 {
         void SetUp();
         void TearDown();
 
-        ALLOCATOR_DESC CreateBasicAllocatorDesc(bool isPrefetchAllowed = false) const;
+        ALLOCATOR_DESC CreateBasicAllocatorDesc() const;
 
         static D3D12_RESOURCE_DESC CreateBasicBufferDesc(uint64_t width, uint64_t alignment = 0);
 
diff --git a/src/tests/capture_replay_tests/D3D12EventTraceReplay.cpp b/src/tests/capture_replay_tests/D3D12EventTraceReplay.cpp
index 285ad707e..b0d33b1e5 100644
--- a/src/tests/capture_replay_tests/D3D12EventTraceReplay.cpp
+++ b/src/tests/capture_replay_tests/D3D12EventTraceReplay.cpp
@@ -375,9 +375,12 @@ class D3D12EventTraceReplay : public D3D12TestBase, public CaptureReplayTestWith
                         const Json::Value& snapshot = event["args"]["snapshot"];
                         ASSERT_FALSE(snapshot.empty());
 
+                        ALLOCATOR_DESC allocatorDesc = CreateBasicAllocatorDesc();
+                        if (!envParams.IsPrefetchAllowed) {
+                            allocatorDesc.Flags |= ALLOCATOR_FLAG_DISABLE_MEMORY_PREFETCH;
+                        }
+
                         // Apply profile (if specified).
-                        ALLOCATOR_DESC allocatorDesc =
-                            CreateBasicAllocatorDesc(envParams.IsPrefetchAllowed);
                         if (envParams.AllocatorProfile ==
                             AllocatorProfile::ALLOCATOR_PROFILE_CAPTURED) {
                             allocatorDesc.Flags |=
diff --git a/src/tests/end2end/D3D12ResourceAllocatorTests.cpp b/src/tests/end2end/D3D12ResourceAllocatorTests.cpp
index 0a914d169..7bf9a5f7e 100644
--- a/src/tests/end2end/D3D12ResourceAllocatorTests.cpp
+++ b/src/tests/end2end/D3D12ResourceAllocatorTests.cpp
@@ -58,6 +58,21 @@ class D3D12ResourceAllocatorTests : public D3D12TestBase, public ::testing::Test
     void TearDown() override {
         D3D12TestBase::TearDown();
     }
+
+    // Configures allocator for testing allocation in a controlled and predictable
+    // fashion.
+    ALLOCATOR_DESC CreateBasicAllocatorDesc() const {
+        ALLOCATOR_DESC desc = D3D12TestBase::CreateBasicAllocatorDesc();
+
+        // Pre-fetching is enabled by default. However for testing purposes, pre-fetching changes
+        // expectations that check GPU memory usage and needs to be tested in isolation.
+        desc.Flags |= ALLOCATOR_FLAG_DISABLE_MEMORY_PREFETCH;
+
+        // Make sure leak detection is always enabled.
+        desc.Flags |= gpgmm::d3d12::ALLOCATOR_FLAG_NEVER_LEAK_MEMORY;
+
+        return desc;
+    }
 };
 
 TEST_F(D3D12ResourceAllocatorTests, CreateAllocator) {
@@ -1520,9 +1535,13 @@ TEST_F(D3D12ResourceAllocatorTests, CreateBufferWithLimitedFragmentation) {
 
 // Creates a bunch of small buffers using the smallest size allowed with pre-fetching enabled.
 TEST_F(D3D12ResourceAllocatorTests, CreateBufferManyPrefetch) {
+    // Prefetching is explicitly disabled but otherwise allowed, re-enable it by clearing the
+    // disable flag.
+    ALLOCATOR_DESC allocatorDesc = CreateBasicAllocatorDesc();
+    allocatorDesc.Flags ^= ALLOCATOR_FLAG_DISABLE_MEMORY_PREFETCH;
+
     ComPtr<ResourceAllocator> resourceAllocator;
-    ASSERT_SUCCEEDED(ResourceAllocator::CreateAllocator(
-        CreateBasicAllocatorDesc(/*allowPrefetch*/ true), &resourceAllocator));
+    ASSERT_SUCCEEDED(ResourceAllocator::CreateAllocator(allocatorDesc, &resourceAllocator));
     ASSERT_NE(resourceAllocator, nullptr);
 
     constexpr uint64_t kNumOfBuffers = 1000u;