From f7ece24b83d4c618960457fea9adc0400167dab6 Mon Sep 17 00:00:00 2001
From: Marcin Lubimow <marcin@campussociety.com>
Date: Thu, 7 Sep 2017 11:14:10 +0100
Subject: [PATCH] Added possibility to use custom schema generator

---
 rest_framework_swagger/views.py | 10 ++++++++--
 tests/test_views.py             | 27 ++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/rest_framework_swagger/views.py b/rest_framework_swagger/views.py
index 851f5102..c6e16e8b 100644
--- a/rest_framework_swagger/views.py
+++ b/rest_framework_swagger/views.py
@@ -8,7 +8,8 @@
 from . import renderers
 
 
-def get_swagger_view(title=None, url=None, patterns=None, urlconf=None):
+def get_swagger_view(title=None, url=None, patterns=None, urlconf=None,
+                     schema_generator_cls=None):
     """
     Returns schema view which renders Swagger/OpenAPI.
     """
@@ -23,7 +24,12 @@ class SwaggerSchemaView(APIView):
         ]
 
         def get(self, request):
-            generator = SchemaGenerator(
+            if schema_generator_cls is None:
+                generator_cls = SchemaGenerator
+            else:
+                generator_cls = schema_generator_cls
+
+            generator = generator_cls(
                 title=title,
                 url=url,
                 patterns=patterns,
diff --git a/tests/test_views.py b/tests/test_views.py
index 7d03c95e..752e5a23 100644
--- a/tests/test_views.py
+++ b/tests/test_views.py
@@ -7,7 +7,7 @@
 from rest_framework_swagger import renderers
 from rest_framework_swagger.views import get_swagger_view
 
-from .compat.mock import patch
+from .compat.mock import patch, MagicMock
 
 
 class TestGetSwaggerView(TestCase):
@@ -16,6 +16,31 @@ def setUp(self):
         self.factory = APIRequestFactory()
         self.view_class = self.sut().cls
 
+    def test_custom_schema_generator_used(self):
+        title = 'Vandelay'
+        url = 'https://github.com/marcgibbons/django-rest-swagger'
+        urlconf = 'fizz'
+        patterns = []
+
+        schema_generator = MagicMock()
+        view = self.sut(
+            title=title,
+            url=url,
+            patterns=patterns,
+            urlconf=urlconf,
+            schema_generator_cls=schema_generator
+        )
+
+        request = self.factory.get('/')
+        view(request=request)
+
+        schema_generator.assert_called_once_with(
+            title=title,
+            url=url,
+            patterns=patterns,
+            urlconf=urlconf
+        )
+
     def test_title_and_urlpassed_to_schema_generator(self):
         title = 'Vandelay'
         url = 'https://github.com/marcgibbons/django-rest-swagger'