1 #!/usr/bin/env python 2 # Copyright 2013 Google Inc. All Rights Reserved. 3 # 4 # Licensed under the Apache License, Version 2.0 (the "License"); 5 # you may not use this file except in compliance with the License. 6 # You may obtain a copy of the License at 7 # 8 # http://www.apache.org/licenses/LICENSE-2.0 9 # 10 # Unless required by applicable law or agreed to in writing, software 11 # distributed under the License is distributed on an "AS IS" BASIS, 12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 # See the License for the specific language governing permissions and 14 # limitations under the License. 15 16 import script_injector 17 import unittest 18 19 20 LONG_COMMENT = '<!--%s-->' % ('comment,' * 200) 21 SCRIPT_TO_INJECT = 'var flag = 0;' 22 EXPECTED_SCRIPT = '<script>%s</script>' % SCRIPT_TO_INJECT 23 TEXT_HTML = 'text/html' 24 TEXT_CSS = 'text/css' 25 APPLICATION = 'application/javascript' 26 27 TEMPLATE_HEAD = '<!doctype html><html><head>%s</head><body></body></html>' 28 TEMPLATE_HTML = '<!doctype html><html>%s<body></body></html>' 29 TEMPLATE_DOCTYPE = '<!doctype html>%s<body></body>' 30 TEMPLATE_RAW = '%s<body></body>' 31 TEMPLATE_COMMENT = '%s<!doctype html>%s<html>%s<head>%s</head></html>' 32 33 34 class ScriptInjectorTest(unittest.TestCase): 35 36 def test_unsupported_content_type(self): 37 source = 'abc' 38 # CSS. 39 new_source, already_injected = script_injector.InjectScript( 40 source, TEXT_CSS, SCRIPT_TO_INJECT) 41 self.assertEqual(new_source, source) 42 self.assertFalse(already_injected) 43 # Javascript. 44 new_source, already_injected = script_injector.InjectScript( 45 source, APPLICATION, SCRIPT_TO_INJECT) 46 self.assertEqual(new_source, source) 47 self.assertFalse(already_injected) 48 49 def test_empty_content_as_already_injected(self): 50 source, already_injected = script_injector.InjectScript( 51 '', TEXT_HTML, SCRIPT_TO_INJECT) 52 self.assertEqual(source, '') 53 self.assertTrue(already_injected) 54 55 def test_already_injected(self): 56 source, already_injected = script_injector.InjectScript( 57 TEMPLATE_HEAD % EXPECTED_SCRIPT, TEXT_HTML, SCRIPT_TO_INJECT) 58 self.assertEqual(source, TEMPLATE_HEAD % EXPECTED_SCRIPT) 59 self.assertTrue(already_injected) 60 61 def _assert_successful_injection(self, template): 62 source, already_injected = script_injector.InjectScript( 63 template % '', TEXT_HTML, SCRIPT_TO_INJECT) 64 self.assertEqual(source, template % EXPECTED_SCRIPT) 65 self.assertFalse(already_injected) 66 67 def test_normal(self): 68 self._assert_successful_injection(TEMPLATE_HEAD) 69 70 def test_no_head_tag(self): 71 self._assert_successful_injection(TEMPLATE_HTML) 72 73 def test_no_head_and_html_tag(self): 74 self._assert_successful_injection(TEMPLATE_DOCTYPE) 75 76 def test_no_head_html_and_doctype_tag(self): 77 self._assert_successful_injection(TEMPLATE_RAW) 78 79 def _assert_successful_injection_with_comment(self, before_doctype, 80 after_doctype, after_html): 81 source, already_injected = script_injector.InjectScript( 82 TEMPLATE_COMMENT % (before_doctype, after_doctype, after_html, ''), 83 TEXT_HTML, SCRIPT_TO_INJECT) 84 expected_source = TEMPLATE_COMMENT % (before_doctype, after_doctype, 85 after_html, EXPECTED_SCRIPT) 86 self.assertEqual(source, expected_source) 87 self.assertFalse(already_injected) 88 89 def test_comment_before_doctype(self): 90 self._assert_successful_injection_with_comment(LONG_COMMENT, '', '') 91 92 def test_comment_after_doctype(self): 93 self._assert_successful_injection_with_comment('', LONG_COMMENT, '') 94 95 def test_comment_after_html(self): 96 self._assert_successful_injection_with_comment('', '', LONG_COMMENT) 97 98 def test_all_comments(self): 99 self._assert_successful_injection_with_comment( 100 LONG_COMMENT, LONG_COMMENT, LONG_COMMENT) 101 102 103 if __name__ == '__main__': 104 unittest.main() 105
