from unittest import TestCase

import sanitize


class TestSanitizeHTML(TestCase):
    def _html(self, html_source, expected_data, base_uri=None, add_nofollow=False):
        """
        :type html_source: str
        :type expected_data: str
        :type base_uri: str
        :type add_nofollow: bool
        """
        self.assertEqual(
            sanitize.HTML(
                htmlSource=html_source,
                baseuri=base_uri,
                addnofollow=add_nofollow
            ),
            expected_data
        )
    
    def test_basics(self):
        self._html("", "")
        self._html("hello", "hello")
    
    def test_balancing_tags(self):
        self._html("<b>hello", "<b>hello</b>")
        self._html("hello<b>", "hello<b></b>")
        self._html("hello</b>", "hello")
        self._html("hello<b/>", "hello<b></b>")
        self._html("<b><b><b>hello", "<b><b><b>hello</b></b></b>")
        self._html("</b><b>", "<b></b>")
    
    def test_trailing_slashes(self):
        self._html('<img>', '<img />')
        self._html('<img/>', '<img />')
        self._html('<b/></b>', '<b></b>')
    
    def test_balancing_angle_brakets(self):
        self._html('<img src="foo"', '')
        self._html('b>', 'b>')
        self._html('<img src="foo"/', '')
        self._html('>', '>')
        self._html('foo<b', 'foo')
        self._html('b>foo', 'b>foo')
        self._html('><b', '>')
        self._html('b><', 'b>')
        self._html('><b>', '><b></b>')

@@ -1476,16 +1476,322 @@
 %3E', '%3E%3Cb%3E%3C/b%3E')%0A
+    %0A    def test_attributes(self):%0A        self._html('%3Cimg src=foo%3E', '%3Cimg src=%22foo%22 /%3E')%0A        self._html('%3Cimg asrc=foo%3E', '%3Cimg /%3E')%0A        self._html('%3Cimg src=test test%3E', '%3Cimg src=%22test%22 /%3E')%0A        self._html('%3Cinput type=%22checkbox%22 checked%3E', '%3Cinput type=%22checkbox%22 checked=%22checked%22 /%3E')%0A

Add test_attributes