From d9f55901431c2c06eeffd1ed79b2c2bf117ecf76 Mon Sep 17 00:00:00 2001
From: Dylan Young <dylanyoungmeijer@gmail.com>
Date: Sun, 8 Mar 2020 23:13:29 -0300
Subject: [PATCH] Update tests to use new multi-policy format

---
 csp/tests/test_decorators.py | 20 ++++++++++++++------
 csp/tests/test_middleware.py |  8 ++++++--
 csp/tests/test_utils.py      | 19 +++++++++++++------
 3 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/csp/tests/test_decorators.py b/csp/tests/test_decorators.py
index 282e051..b4512fe 100644
--- a/csp/tests/test_decorators.py
+++ b/csp/tests/test_decorators.py
@@ -32,7 +32,7 @@ def view_without_decorator(request):
     def view_with_decorator(request):
         return HttpResponse()
     response = view_with_decorator(REQUEST)
-    assert response._csp_update == {'img-src': 'bar.com'}
+    assert dict(response._csp_update) == {'default': {'img-src': ['bar.com']}}
     mw.process_response(REQUEST, response)
     policy_list = sorted(response['Content-Security-Policy'].split("; "))
     assert policy_list == ["default-src 'self'", "img-src foo.com bar.com"]
@@ -56,7 +56,7 @@ def view_without_decorator(request):
     def view_with_decorator(request):
         return HttpResponse()
     response = view_with_decorator(REQUEST)
-    assert response._csp_replace == {'img-src': 'bar.com'}
+    assert dict(response._csp_replace) == {'default': {'img-src': ['bar.com']}}
     mw.process_response(REQUEST, response)
     policy_list = sorted(response['Content-Security-Policy'].split("; "))
     assert policy_list == ["default-src 'self'", "img-src bar.com"]
@@ -87,8 +87,12 @@ def view_without_decorator(request):
     def view_with_decorator(request):
         return HttpResponse()
     response = view_with_decorator(REQUEST)
-    assert response._csp_config == \
-        {'img-src': ['foo.com'], 'font-src': ['bar.com']}
+    assert response._csp_config == {
+        policy_names.last_policy_name: {
+            'img-src': ['foo.com'],
+            'font-src': ['bar.com'],
+        }
+    }
     mw.process_response(REQUEST, response)
     policy_list = sorted(response['Content-Security-Policy'].split("; "))
     assert policy_list == ["font-src bar.com", "img-src foo.com"]
@@ -105,8 +109,12 @@ def test_csp_string_values():
     def view_with_decorator(request):
         return HttpResponse()
     response = view_with_decorator(REQUEST)
-    assert response._csp_config == \
-        {'img-src': ['foo.com'], 'font-src': ['bar.com']}
+    assert dict(response._csp_config) == {
+        policy_names.last_policy_name: {
+            'img-src': ['foo.com'],
+            'font-src': ['bar.com'],
+        }
+    }
     mw.process_response(REQUEST, response)
     policy_list = sorted(response['Content-Security-Policy'].split("; "))
     assert policy_list == ["font-src bar.com", "img-src foo.com"]
diff --git a/csp/tests/test_middleware.py b/csp/tests/test_middleware.py
index f9413ce..cf03250 100644
--- a/csp/tests/test_middleware.py
+++ b/csp/tests/test_middleware.py
@@ -57,7 +57,9 @@ def test_dont_replace():
 def test_use_config():
     request = rf.get('/')
     response = HttpResponse()
-    response._csp_config = {'default-src': ['example.com']}
+    response._csp_config = {'default': {
+        'default-src': ['example.com'],
+    }}
     mw.process_response(request, response)
     assert response[HEADER] == 'default-src example.com'
 
@@ -65,7 +67,9 @@ def test_use_config():
 def test_use_update():
     request = rf.get('/')
     response = HttpResponse()
-    response._csp_update = {'default-src': ['example.com']}
+    response._csp_update = {'default': {
+        'default-src': ['example.com']
+    }}
     mw.process_response(request, response)
     assert response[HEADER] == "default-src 'self' example.com"
 
diff --git a/csp/tests/test_utils.py b/csp/tests/test_utils.py
index 4bb0659..5996fea 100644
--- a/csp/tests/test_utils.py
+++ b/csp/tests/test_utils.py
@@ -9,15 +9,22 @@
 from csp.utils import build_policy
 
 
-def policy_eq(a, b, msg='%r != %r'):
-    parts_a = sorted(a.split('; '))
-    parts_b = sorted(b.split('; '))
-    assert parts_a == parts_b, msg % (a, b)
+def policy_eq(a, b, msg='%r != %r', report_only=False):
+    if not isinstance(a, list):
+        a = [(a, report_only)]
+    if not isinstance(a, list):
+        b = [(b, report_only)]
+
+    for csp_a, csp_b in zip(a, b):
+        parts_a = sorted(csp_a[0].split('; '))
+        parts_b = sorted(csp_b[0].split('; '))
+        assert csp_a[1] == csp_b[1]
+        assert parts_a == parts_b, msg % (a, b)
 
 
 def test_empty_policy():
     policy = build_policy()
-    assert "default-src 'self'" == policy
+    assert [("default-src 'self'", False)] == policy
 
 
 def literal(s):
@@ -30,7 +37,7 @@ def literal(s):
 @override_settings(CSP_DEFAULT_SRC=['example.com', 'example2.com'])
 def test_default_src():
     policy = build_policy()
-    assert 'default-src example.com example2.com' == policy
+    assert [('default-src example.com example2.com', False)] == policy
 
 
 @override_settings(CSP_SCRIPT_SRC=['example.com'])