diff --git a/BUILD.gn b/BUILD.gn
index aa45294..2d90fca 100644
--- a/BUILD.gn
+++ b/BUILD.gn
@@ -474,9 +474,6 @@
   if (is_android) {
     use_raw_android_executable = true
   }
-  if (pdf_use_partition_alloc) {
-    deps += [ "//base/allocator/partition_allocator:partition_alloc" ]
-  }
   if (pdf_enable_v8) {
     configs += [ "//v8:external_startup_data" ]
     deps += [
diff --git a/samples/BUILD.gn b/samples/BUILD.gn
index 577d685..981750b 100644
--- a/samples/BUILD.gn
+++ b/samples/BUILD.gn
@@ -21,7 +21,9 @@
   if (is_asan) {
     defines += [ "PDF_ENABLE_ASAN" ]
   }
-
+  if (pdf_use_partition_alloc) {
+    defines += [ "PDF_USE_PARTITION_ALLOC" ]
+  }
   if (enable_callgrind) {
     defines += [ "ENABLE_CALLGRIND" ]
   }
diff --git a/samples/pdfium_test.cc b/samples/pdfium_test.cc
index 1dfc70c..282c539 100644
--- a/samples/pdfium_test.cc
+++ b/samples/pdfium_test.cc
@@ -62,6 +62,10 @@
 #include <valgrind/callgrind.h>
 #endif  // ENABLE_CALLGRIND
 
+#if defined(PDF_USE_PARTITION_ALLOC)
+#include "testing/allocator_shim_config.h"
+#endif
+
 #ifdef PDF_ENABLE_SKIA
 #include "third_party/skia/include/core/SkCanvas.h"           // nogncheck
 #include "third_party/skia/include/core/SkColor.h"            // nogncheck
@@ -1832,6 +1836,10 @@
 }  // namespace
 
 int main(int argc, const char* argv[]) {
+#if defined(PDF_USE_PARTITION_ALLOC)
+  pdfium::ConfigurePartitionAllocShimPartitionForTest();
+#endif
+
   SetUpErrorHandling();
   setlocale(LC_CTYPE, "en_US.UTF-8");  // For printf() of high-characters.
 
diff --git a/testing/BUILD.gn b/testing/BUILD.gn
index a20955a..309e6c3 100644
--- a/testing/BUILD.gn
+++ b/testing/BUILD.gn
@@ -51,6 +51,13 @@
     "../:pdfium_noshorten_config",
   ]
   visibility = [ "../*" ]
+  if (pdf_use_partition_alloc) {
+    sources += [
+      "allocator_shim_config.cpp",
+      "allocator_shim_config.h",
+    ]
+    deps += [ "//base/allocator/partition_allocator:partition_alloc" ]
+  }
   if (pdf_enable_v8) {
     sources += [
       "v8_initializer.cpp",
diff --git a/testing/allocator_shim_config.cpp b/testing/allocator_shim_config.cpp
new file mode 100644
index 0000000..8f0380d
--- /dev/null
+++ b/testing/allocator_shim_config.cpp
@@ -0,0 +1,25 @@
+// Copyright 2023 The PDFium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#include "testing/allocator_shim_config.h"
+
+#include "base/allocator/partition_allocator/partition_alloc_buildflags.h"
+#include "base/allocator/partition_allocator/shim/allocator_shim.h"
+
+namespace pdfium {
+
+void ConfigurePartitionAllocShimPartitionForTest() {
+#if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
+#if BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
+  allocator_shim::ConfigurePartitions(
+      allocator_shim::EnableBrp(true),
+      allocator_shim::EnableMemoryTagging(false),
+      allocator_shim::SplitMainPartition(true),
+      allocator_shim::UseDedicatedAlignedPartition(true), 0,
+      allocator_shim::BucketDistribution::kNeutral);
+#endif  // BUILDFLAG(ENABLE_BACKUP_REF_PTR_SUPPORT)
+#endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
+}
+
+}  // namespace pdfium
diff --git a/testing/allocator_shim_config.h b/testing/allocator_shim_config.h
new file mode 100644
index 0000000..cd14484
--- /dev/null
+++ b/testing/allocator_shim_config.h
@@ -0,0 +1,18 @@
+// Copyright 2023 The PDFium Authors
+// Use of this source code is governed by a BSD-style license that can be
+// found in the LICENSE file.
+
+#ifndef TESTING_ALLOCATOR_SHIM_CONFIG_H_
+#define TESTING_ALLOCATOR_SHIM_CONFIG_H_
+
+#if !defined(PDF_USE_PARTITION_ALLOC)
+#error "Included under the wrong build options"
+#endif
+
+namespace pdfium {
+
+void ConfigurePartitionAllocShimPartitionForTest();
+
+}  // namespace pdfium
+
+#endif  // TESTING_ALLOCATOR_SHIM_CONFIG_H_
diff --git a/testing/embedder_test_main.cpp b/testing/embedder_test_main.cpp
index 53049b6..e089b1a 100644
--- a/testing/embedder_test_main.cpp
+++ b/testing/embedder_test_main.cpp
@@ -12,10 +12,18 @@
 #include "testing/v8_test_environment.h"
 #endif  // PDF_ENABLE_V8
 
+#if defined(PDF_USE_PARTITION_ALLOC)
+#include "testing/allocator_shim_config.h"
+#endif
+
 // Can't use gtest-provided main since we need to create our own
 // testing environment which needs the executable path in order to
 // find the external V8 binary data files.
 int main(int argc, char** argv) {
+#if defined(PDF_USE_PARTITION_ALLOC)
+  pdfium::ConfigurePartitionAllocShimPartitionForTest();
+#endif
+
   FX_InitializeMemoryAllocators();
 
 #ifdef PDF_ENABLE_V8
diff --git a/testing/unit_test_main.cpp b/testing/unit_test_main.cpp
index 30d4203..28476b0 100644
--- a/testing/unit_test_main.cpp
+++ b/testing/unit_test_main.cpp
@@ -8,8 +8,7 @@
 #include "testing/pdf_test_environment.h"
 
 #if defined(PDF_USE_PARTITION_ALLOC)
-#include "base/allocator/partition_allocator/partition_alloc_buildflags.h"
-#include "base/allocator/partition_allocator/shim/allocator_shim.h"
+#include "testing/allocator_shim_config.h"
 #endif
 
 #ifdef PDF_ENABLE_V8
@@ -23,14 +22,7 @@
 // alloc before invoking any test, and add test environments.
 int main(int argc, char** argv) {
 #if defined(PDF_USE_PARTITION_ALLOC)
-#if BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
-  allocator_shim::ConfigurePartitions(
-      allocator_shim::EnableBrp(true),
-      allocator_shim::EnableMemoryTagging(false),
-      allocator_shim::SplitMainPartition(true),
-      allocator_shim::UseDedicatedAlignedPartition(true), 0,
-      allocator_shim::BucketDistribution::kNeutral);
-#endif  // BUILDFLAG(USE_PARTITION_ALLOC_AS_MALLOC)
+  pdfium::ConfigurePartitionAllocShimPartitionForTest();
 #endif  // defined(PDF_USE_PARTITION_ALLOC)
 
   FX_InitializeMemoryAllocators();