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 16import script_injector 17import unittest 18 19 20LONG_COMMENT = '<!--%s-->' % ('comment,' * 200) 21SCRIPT_TO_INJECT = 'var flag = 0;' 22EXPECTED_SCRIPT = '<script>%s</script>' % SCRIPT_TO_INJECT 23TEXT_HTML = 'text/html' 24TEXT_CSS = 'text/css' 25APPLICATION = 'application/javascript' 26 27TEMPLATE_HEAD = '<!doctype html><html><head>%s</head><body></body></html>' 28TEMPLATE_HTML = '<!doctype html><html>%s<body></body></html>' 29TEMPLATE_DOCTYPE = '<!doctype html>%s<body></body>' 30TEMPLATE_RAW = '%s<body></body>' 31TEMPLATE_COMMENT = '%s<!doctype html>%s<html>%s<head>%s</head></html>' 32 33 34class 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 103if __name__ == '__main__': 104 unittest.main() 105