From 5558af25272b3863f51ef6ed125ab83a19baca31 Mon Sep 17 00:00:00 2001
From: Chris Hill-Scott <me@quis.cc>
Date: Thu, 6 Jan 2022 11:06:55 +0000
Subject: [PATCH] Refactor for reuse

---
 app/main/validators.py                       | 23 +++++++++++++++----
 tests/app/main/views/test_letter_branding.py | 24 +++++++++++++-------
 2 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/app/main/validators.py b/app/main/validators.py
index ec46113ac..cda57b3a5 100644
--- a/app/main/validators.py
+++ b/app/main/validators.py
@@ -1,4 +1,5 @@
 import re
+from abc import ABC, abstractmethod
 
 from notifications_utils.field import Field
 from notifications_utils.formatters import formatted_list
@@ -77,18 +78,30 @@ class NoCommasInPlaceHolders:
             raise ValidationError(self.message)
 
 
-class NoEmbeddedImagesInSVG:
+class NoElementInSVG(ABC):
 
-    def __init__(self, message='This SVG has an embedded raster image in it and will not render well'):
-        self.message = message
+    @property
+    @abstractmethod
+    def element(self):
+        pass
+
+    @property
+    @abstractmethod
+    def message(self):
+        pass
 
     def __call__(self, form, field):
-        is_image_embedded = '<image' in field.data.stream.read().decode("utf-8")
+        svg_contents = field.data.stream.read().decode("utf-8")
         field.data.stream.seek(0)
-        if is_image_embedded:
+        if f'<{self.element}' in svg_contents.lower():
             raise ValidationError(self.message)
 
 
+class NoEmbeddedImagesInSVG(NoElementInSVG):
+    element = 'image'
+    message = 'This SVG has an embedded raster image in it and will not render well'
+
+
 class OnlySMSCharacters:
 
     def __init__(self, *args, template_type, **kwargs):
diff --git a/tests/app/main/views/test_letter_branding.py b/tests/app/main/views/test_letter_branding.py
index a2bd8f992..4a5bdcbe6 100644
--- a/tests/app/main/views/test_letter_branding.py
+++ b/tests/app/main/views/test_letter_branding.py
@@ -2,6 +2,7 @@ from io import BytesIO
 from unittest.mock import Mock, call
 from uuid import UUID
 
+import pytest
 from botocore.exceptions import ClientError as BotoClientError
 from bs4 import BeautifulSoup
 from flask import url_for
@@ -348,10 +349,21 @@ def test_create_letter_branding_when_uploading_valid_file(
     assert mock_delete_temp_files.called is False
 
 
-def test_create_letter_branding_fails_validation_when_uploading_SVG_with_embedded_image(
+@pytest.mark.parametrize('svg_contents, expected_error', (
+    (
+        '''
+            <svg height="100" width="100">
+            <image href="someurlgoeshere" x="0" y="0" height="100" width="100"></image></svg>
+        ''',
+        'This SVG has an embedded raster image in it and will not render well',
+    ),
+))
+def test_create_letter_branding_fails_validation_when_uploading_SVG_with_bad_element(
     mocker,
     platform_admin_client,
-    fake_uuid
+    fake_uuid,
+    svg_contents,
+    expected_error,
 ):
     filename = 'test.svg'
 
@@ -359,18 +371,14 @@ def test_create_letter_branding_fails_validation_when_uploading_SVG_with_embedde
 
     response = platform_admin_client.post(
         url_for('.create_letter_branding'),
-        data={'file': (BytesIO("""
-            <svg height="100" width="100">
-            <image href="someurlgoeshere" x="0" y="0" height="100" width="100"></image></svg>
-        """.encode('utf-8')), filename)},
+        data={'file': (BytesIO(svg_contents.encode('utf-8')), filename)},
         follow_redirects=True,
     )
 
     page = BeautifulSoup(response.data.decode('utf-8'), 'html.parser')
 
     assert normalize_spaces(page.find('h1').text) == "Add letter branding"
-    message = 'This SVG has an embedded raster image in it and will not render well'
-    assert normalize_spaces(page.find("span", {"class": "error-message"}).text) == message
+    assert normalize_spaces(page.select_one(".error-message").text) == expected_error
 
     assert page.findAll('div', {'id': 'logo-img'}) == []