From d575b28d1b47fb50bb181fe211a9664599a5ac46 Mon Sep 17 00:00:00 2001
From: Daniel Sotirhos <dsotirho@ucsc.edu>
Date: Wed, 5 Apr 2023 17:28:08 -0700
Subject: [PATCH] fixup! Fix: Frame Injection (azul-private#12)

---
 test/test_content_type.py | 265 +++++++++++++-------------------------
 1 file changed, 92 insertions(+), 173 deletions(-)

diff --git a/test/test_content_type.py b/test/test_content_type.py
index 1597488de8..5cfa1ede47 100644
--- a/test/test_content_type.py
+++ b/test/test_content_type.py
@@ -2,6 +2,7 @@
 from typing import (
     Any,
     Callable,
+    Tuple,
 )
 from unittest import (
     mock,
@@ -18,30 +19,17 @@
 from app_test_case import (
     LocalAppTestCase,
 )
-from azul import (
-    CatalogName,
-    JSON,
-    config,
+from azul_test_case import (
+    DCP2TestCase,
 )
 
 
-class TestContentType(LocalAppTestCase):
+class TestContentType(LocalAppTestCase, DCP2TestCase):
 
     @classmethod
     def lambda_name(cls) -> str:
         return 'service'
 
-    @classmethod
-    def catalog_config(cls) -> dict[CatalogName, config.Catalog]:
-        return {
-            cls.catalog: config.Catalog(name=cls.catalog,
-                                        atlas='hca',
-                                        internal=False,
-                                        plugins=dict(metadata=config.Catalog.Plugin(name='hca'),
-                                                     repository=config.Catalog.Plugin(name='dss')),
-                                        sources=set())
-        }
-
     @classmethod
     def setUpClass(cls):
         super().setUpClass()
@@ -52,7 +40,7 @@ def route():
 
     def _replace_handler(self, handler: Callable[[None], Any]):
         """
-        Replace the current handler for route `/test` with the provided
+        Replace the current handler for route `/test` with the provided handler
         """
         route = '/test'
         app = self.__class__.app_module.app
@@ -67,7 +55,7 @@ def _replace_handler(self, handler: Callable[[None], Any]):
     def _shrink_traceback(self, s: str) -> str:
         """
         Return a modified version of the given traceback. The first and last
-        lines are kept, and all lines inbetween are replaced with a single line
+        lines are kept, and everything inbetween is replaced with a single line
         of '...'.
         """
         if s.startswith('Traceback'):
@@ -84,30 +72,32 @@ def _shrink_traceback(self, s: str) -> str:
             s = re.sub(pattern, r'\1\2...\\n\3\4', s)
         return s
 
-    def _test_route(self, handler_fn: Callable[[None], Any], expected_fn: Callable[[bool], JSON]):
+    def _test_route(self,
+                    handler_fn: Callable[[None], Any],
+                    expected_fn: Callable[[bool, bool], Tuple[str, str]]
+                    ):
         """
-        Verify response values against expected for requests made with various
+        Verify the response against expected for requests made with various
         types of `accept` header values.
         """
         self._replace_handler(handler_fn)
-        for debug in (0, 1):
-            with mock.patch.object(Chalice, 'debug', debug):
-                expected = expected_fn(bool(debug))
-                for accept, expect_escaped in [
+        for debug in (False, True):
+            with mock.patch.object(Chalice, 'debug', 1 if debug else 0):
+                for accept, expect_wrapped in [
                     (None, False),
                     ('*/*', False),
                     ('*/*,text/html', False),
                     ('text/html', True),
                     ('text/html,*/*', True),
+                    ('*/*;q=0.9,text/html', True),
+                    ('text/html;q=0.9,*/*;q=1.0', False),
                 ]:
                     with self.subTest(debug=debug, accept=accept):
                         url = self.base_url.set(path=('test',))
                         headers = {'accept': accept}
                         response = requests.get(url, headers=headers)
                         response_text = self._shrink_traceback(response.text)
-                        type_ = 'escaped' if expect_escaped else 'unescaped'
-                        expected_text = expected[type_]['text']
-                        expected_content_type = expected[type_]['content-type']
+                        expected_text, expected_content_type = expected_fn(debug, expect_wrapped)
                         self.assertEqual(expected_text, response_text)
                         self.assertEqual(expected_content_type, response.headers['Content-Type'])
 
@@ -117,17 +107,10 @@ def route():
             return '<script />'
 
         # noinspection PyUnusedLocal
-        def expected(debug: bool) -> JSON:
-            return {
-                'unescaped': {
-                    'text': '<script />',
-                    'content-type': 'application/json'
-                },
-                'escaped': {
-                    'text': '<script />',
-                    'content-type': 'application/json'
-                }
-            }
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
+            text = '<script />'
+            content_type = 'application/json'
+            return text, content_type
 
         self._test_route(route, expected)
 
@@ -137,17 +120,10 @@ def route():
             return {'<script />': '<iframe></iframe>'}
 
         # noinspection PyUnusedLocal
-        def expected(debug: bool) -> JSON:
-            return {
-                'unescaped': {
-                    'text': '{"<script />":"<iframe></iframe>"}',
-                    'content-type': 'application/json'
-                },
-                'escaped': {
-                    'text': '{"<script />":"<iframe></iframe>"}',
-                    'content-type': 'application/json'
-                }
-            }
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
+            text = '{"<script />":"<iframe></iframe>"}'
+            content_type = 'application/json'
+            return text, content_type
 
         self._test_route(route, expected)
 
@@ -157,17 +133,10 @@ def route():
             return Response(status_code=200, body='<script />')
 
         # noinspection PyUnusedLocal
-        def expected(debug: bool) -> JSON:
-            return {
-                'unescaped': {
-                    'text': '<script />',
-                    'content-type': 'application/json'
-                },
-                'escaped': {
-                    'text': '<script />',
-                    'content-type': 'application/json'
-                }
-            }
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
+            text = '<script />'
+            content_type = 'application/json'
+            return text, content_type
 
         self._test_route(route, expected)
 
@@ -179,22 +148,19 @@ def route():
                             body='<script />')
 
         # noinspection PyUnusedLocal
-        def expected(debug: bool) -> JSON:
-            return {
-                'unescaped': {
-                    'text': '<script />',
-                    'content-type': 'text/plain'
-                },
-                'escaped': {
-                    'text': (
-                        '<html><body>'
-                        '<h1>Status 200</h1>'
-                        '<pre>"&lt;script /&gt;"</pre>'
-                        '</body></html>'
-                    ),
-                    'content-type': 'text/html'
-                }
-            }
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
+            if wrapped:
+                text = (
+                    '<html><body>'
+                    '<h1>Status 200</h1>'
+                    '<pre>"&lt;script /&gt;"</pre>'
+                    '</body></html>'
+                )
+                content_type = 'text/html'
+            else:
+                text = '<script />'
+                content_type = 'text/plain'
+            return text, content_type
 
         self._test_route(route, expected)
 
@@ -206,17 +172,10 @@ def route():
                             body='<script />')
 
         # noinspection PyUnusedLocal
-        def expected(debug: bool) -> JSON:
-            return {
-                'unescaped': {
-                    'text': '<script />',
-                    'content-type': 'text/html'
-                },
-                'escaped': {
-                    'text': '<script />',
-                    'content-type': 'text/html'
-                }
-            }
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
+            text = '<script />'
+            content_type = 'text/html'
+            return text, content_type
 
         self._test_route(route, expected)
 
@@ -226,17 +185,10 @@ def route():
             raise NotFoundError('<script />')
 
         # noinspection PyUnusedLocal
-        def expected(debug: bool) -> JSON:
-            return {
-                'unescaped': {
-                    'text': '{"Code":"NotFoundError","Message":"<script />"}',
-                    'content-type': 'application/json'
-                },
-                'escaped': {
-                    'text': '{"Code":"NotFoundError","Message":"<script />"}',
-                    'content-type': 'application/json'
-                }
-            }
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
+            text = '{"Code":"NotFoundError","Message":"<script />"}'
+            content_type = 'application/json'
+            return text, content_type
 
         self._test_route(route, expected)
 
@@ -245,40 +197,22 @@ def test_ChaliceUnhandledError(self):
         def route():
             raise ChaliceUnhandledError('<script />')
 
-        def expected(debug: bool) -> JSON:
-            return {
-                **(
-                    {
-                        'unescaped': {
-                            'text': (
-                                'Traceback (most recent call last):\n'
-                                '...\n'
-                                'chalice.app.ChaliceUnhandledError: <script />'
-                            ),
-                            'content-type': 'text/plain'
-                        },
-                        'escaped': {
-                            'text': (
-                                'Traceback (most recent call last):\n'
-                                '...\n'
-                                'chalice.app.ChaliceUnhandledError: <script />'
-                            ),
-                            'content-type': 'text/plain'
-                        }
-                    } if debug else {
-                        'unescaped': {
-                            'text': '{"Code":"InternalServerError",'
-                                    '"Message":"An internal server error occurred."}',
-                            'content-type': 'application/json'
-                        },
-                        'escaped': {
-                            'text': '{"Code":"InternalServerError",'
-                                    '"Message":"An internal server error occurred."}',
-                            'content-type': 'application/json'
-                        }
-                    }
-                ),
-            }
+        # noinspection PyUnusedLocal
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
+            if debug:
+                text = (
+                    'Traceback (most recent call last):\n'
+                    '...\n'
+                    'chalice.app.ChaliceUnhandledError: <script />'
+                )
+                content_type = 'text/plain'
+            else:
+                text = (
+                    '{"Code":"InternalServerError",'
+                    '"Message":"An internal server error occurred."}'
+                )
+                content_type = 'application/json'
+            return text, content_type
 
         self._test_route(route, expected)
 
@@ -287,48 +221,33 @@ def test_exception(self):
         def route():
             raise Exception('<script />')
 
-        def expected(debug: bool) -> JSON:
+        def expected(debug: bool, wrapped: bool) -> Tuple[str, str]:
             # Chalice's `_unhandled_exception_to_response` returns a
             # stacktrace if debug is enabled
-            return {
-                **(
-                    {
-                        'unescaped': {
-                            'text': (
-                                'Traceback (most recent call last):\n'
-                                '...\n'
-                                'Exception: <script />'
-                            ),
-                            'content-type': 'text/plain'
-                        },
-                        'escaped': {
-                            'text': (
-                                '<html><body>'
-                                '<h1>Status 500</h1>'
-                                '<pre>"Traceback (most recent call last):\\n'
-                                '...\\n'
-                                'Exception: &lt;script /&gt;"</pre>'
-                                '</body></html>'
-                            ),
-                            'content-type': 'text/html'
-                        }
-                    } if debug else {
-                        'unescaped': {
-                            'text': (
-                                '{"Code":"InternalServerError",'
-                                '"Message":"An internal server error occurred."}'
-                            ),
-                            'content-type': 'application/json'
-                        },
-                        'escaped': {
-                            'text': (
-                                '{"Code":"InternalServerError",'
-                                '"Message":"An internal server error occurred."}'
-                            ),
-                            'content-type': 'application/json'
-                        }
-                    }
-                ),
-            }
+            if debug:
+                if wrapped:
+                    text = (
+                        '<html><body>'
+                        '<h1>Status 500</h1>'
+                        '<pre>"Traceback (most recent call last):\\n'
+                        '...\\n'
+                        'Exception: &lt;script /&gt;"</pre>'
+                        '</body></html>'
+                    )
+                    content_type = 'text/html'
+                else:
+                    text = (
+                        'Traceback (most recent call last):\n'
+                        '...\n'
+                        'Exception: <script />'
+                    )
+                    content_type = 'text/plain'
+            else:
+                text = (
+                    '{"Code":"InternalServerError",'
+                    '"Message":"An internal server error occurred."}'
+                )
+                content_type = 'application/json'
+            return text, content_type
 
         self._test_route(route, expected)