From eaa9e65f667f14cb5ac821a7fb2ef2edbc50c3ac Mon Sep 17 00:00:00 2001
From: William Chu <william.chu@uptickhq.com>
Date: Wed, 19 Jun 2024 10:00:11 +1000
Subject: [PATCH] feat: use custom test client class with
 DJANGO_PEV_EXPLAIN_TEST_CLASS

---
 README.md           | 10 +++++++++-
 django_pev/views.py |  3 ++-
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 803ef10..1e0b440 100644
--- a/README.md
+++ b/README.md
@@ -35,6 +35,7 @@ INSTALLED_APPS = [
 ]
 ```
 
+
 # Usage
 
 Wrap some code with the explain context manager. All sql queries are captured
@@ -69,6 +70,13 @@ print(e.slowest.stacktrace)
 pev_response.delete()
 ```
 
+Optionally configure additional settings:
+```python
+# Replace the default test client used during explain with a custom class
+DJANGO_PEV_EXPLAIN_TEST_CLIENT = 'django.test.Client'
+
+```
+
 **How to debug a slow endpoint in production**
 
 If you have access to `python manage.py shell` on the production server;
@@ -96,7 +104,7 @@ print(e.slowest.visualize(title=f"Fetching {url}"))
 - [ ] Add migration to ensure pg_stats_statement_info is correct
 - [ ] Do not crash when its not available
 - [ ] Add explain tab
-= [ ] Add index suggester
+- [ ] Add index suggester
 
 # Disclaimer
 
diff --git a/django_pev/views.py b/django_pev/views.py
index 859d1a4..4b789cf 100644
--- a/django_pev/views.py
+++ b/django_pev/views.py
@@ -8,8 +8,8 @@
 from django.http import HttpRequest, HttpResponseRedirect
 from django.http.response import HttpResponse as HttpResponse
 from django.shortcuts import render
-from django.test import Client as TestClient
 from django.urls import reverse
+from django.utils.module_loading import import_string
 from django.views.generic import FormView, TemplateView
 
 from .utils import ExplainSet, explain, indexes, live_connections, maintenance, queries, space
@@ -116,6 +116,7 @@ class ExplainView(BaseView):
     def get_context_data(self, **kwargs: Any) -> dict[str, Any]:
         ctx = super().get_context_data(**kwargs)
         explain_result = None
+        TestClient = import_string(getattr(settings, "DJANGO_PEV_EXPLAIN_TEST_CLIENT", "django.test.Client"))
 
         if url := self.request.GET.get("url"):
             client = TestClient()