def test_open_browser(self, mock_launch, mock_set_endpoint):
mock_launch.side_effect = mocked_launch
mock_set_endpoint.side_effect = mocked_set_endpoint
browser1 = to_sync(open_browser(True, write_websocket=False))
self.assertTrue(browser1.headless)
self.assertFalse(browser1.autoClose)
self.assertEqual(browser1.args, None)
self.assertEqual(browser1.wsEndpoint, "wsEndpoint")
browser2 = to_sync(open_browser(False, write_websocket=False))
self.assertFalse(browser2.headless)
self.assertFalse(browser2.autoClose)
self.assertEqual(browser2.args, None)
self.assertEqual(browser2.wsEndpoint, "wsEndpoint")
browser3 = to_sync(
open_browser(True,
launch_args=["--no-sandbox"],
write_websocket=False))
self.assertTrue(browser3.headless)
self.assertFalse(browser3.autoClose)
self.assertEqual(browser3.args, ["--no-sandbox"])
self.assertEqual(browser3.wsEndpoint, "wsEndpoint")
with self.assertRaisesRegex(ValueError, "wsEndpoint"):
browser4 = to_sync(open_browser(True,
launch_args=["--no-sandbox"]))
self.assertTrue(browser4.headless)
self.assertFalse(browser4.autoClose)
self.assertEqual(browser4.args, ["--no-sandbox"])
self.assertEqual(browser4.wsEndpoint, "wsEndpoint")
def test_screenshot_protected_page_no_auth(self):
with self.assertRaisesRegex(BadSelector, 'selector unknown: "#godot"'):
to_sync(
generate_bytes_img(
"http://localhost:5000/protected_index", selector="#godot"
)
)
def setUp(self):
to_sync(get_browser(launch_args=["--no-sandbox"]))
self.img_dog = Image.open("tests/server/static/images/aww_dog.jpg")
self.img_dog_change = Image.open(
"tests/server/static/images/aww_dog_change.jpg"
)
self.img_kitten = Image.open("tests/server/static/images/aww_kitten.jpg")
def test_get_browser(self, mock_open_browser, mock_connect,
mock_get_endpoint):
mock_connect.side_effect = mocked_connect
mock_open_browser.side_effect = mocked_open_browser
mock_get_endpoint.return_value = "endpoint"
with self.assertRaisesRegex(ValueError, "connect !"):
to_sync(get_browser())
mock_get_endpoint.return_value = None
with self.assertRaisesRegex(ValueError, "open_browser !"):
to_sync(get_browser())
def test_get_token_from_local_storage(self, mock_post):
mock_post.return_value.status_code = 200
mock_post.return_value.content = '{"token": "hey you"}'
browser = to_sync(
open_browser(True,
write_websocket=False,
launch_args=["--no-sandbox"]))
page = to_sync(browser.newPage())
to_sync(page.goto('http://duckduckgo.com'))
get_token("http://false", {}, True, page)
token = get_local_storage_token(page)
self.assertEqual(token, '{"token": "hey you"}')
browser.close()
def test_screamshot_same_kitten(self):
"""
Takes a screenshot of a given div and compares it to the real image.
"""
to_sync(
generate_bytes_img(
"http://localhost:5000/index.html",
selector="#caterpillar",
path="test_img.jpg",
)
)
img = Image.open("test_img.jpg")
self.assertTrue(_is_same_image(img, self.img_kitten))
def test_screamshot_different_image(self):
"""
Takes a screenshot of a given div and compares it to a different image.
"""
to_sync(
generate_bytes_img(
"http://localhost:5000/index.html",
selector="#godot",
path="test_img.jpg",
)
)
img = Image.open("test_img.jpg")
self.assertEqual(img.size, self.img_dog.size)
self.assertNotAlmostEqual(_rmsd(img, self.img_dog_change), 0, delta=0.05)
def test_screenshot_protected_page_bad_auth(self):
with self.assertRaisesRegex(BadSelector, 'selector unknown: "#godot"'):
to_sync(
generate_bytes_img(
"http://localhost:5000/protected_index",
selector="#godot",
credentials={"token_in_header": True, "token": "xxx"},
)
)
with self.assertRaisesRegex(BadSelector, 'selector unknown: "#godot"'):
to_sync(
generate_bytes_img(
"http://localhost:5000/protected_index",
selector="#godot",
credentials={"token": TOKEN},
)
)
def test_open_headless_browser(self):
"""
Opens a headless browser, checks if it is open and closes it
"""
run(['python3', 'screamshot/browser_manager_script.py', '-o', '-ns'])
endpoint = get_endpoint()
browser = to_sync(connect({'browserWSEndpoint': endpoint}))
self.assertIsInstance(browser, Browser)
run(['python3', 'screamshot/browser_manager_script.py', '-c'])
def test_generate_bytes_img(self, mock_get_browser):
"""
Tests generate_bytes_img
"""
mock_get_browser.side_effect = get_browser
url = "http://fake"
screenshot1 = to_sync(generate_bytes_img(url))
self.assertEqual(screenshot1, "screenshot !")
def main():
parser = ArgumentParser(description=__doc__)
# Mandatory arguments
group = parser.add_mutually_exclusive_group()
group.add_argument("-c", "--close", action="store_true", help="Close the browsers in the \
endpointlist.txt file")
group.add_argument("-o", "--open", action="store_true", help="Open a browser and store its \
websocket endpoint in endpointlist.txt")
# Optionnal argument
parser.add_argument("-g", "--graphic", dest="headless", action="store_false", help="Open the \
browser in graphic mode")
parser.add_argument("-ns", "--no-sandbox", action="store_const", const=["--no-sandbox"],
default="[]", help="Open the browser without sandbox")
args = parser.parse_args()
if args.close:
to_sync(close_browser())
if args.open:
to_sync(open_browser(args.headless, launch_args=args.no_sandbox))
def test_screenshot_protected_page_with_auth_login(self):
img_bytes = BytesIO(
to_sync(
generate_bytes_img(
"http://localhost:5000/protected_index",
selector="#godot",
credentials={"username": "******", "password": "******"},
)
)
)
img = Image.open(img_bytes)
dog_img = Image.open("tests/server/static/images/aww_dog.jpg").convert("RGBA")
kitten_img = Image.open("tests/server/static/images/aww_kitten.jpg").convert(
"RGBA"
)
self.assertTrue(_is_same_image(img, dog_img))
self.assertFalse(_is_same_image(img, kitten_img))
def test_screamshot_same_bytes_write(self):
"""
Takes a screenshot and compares the buffer to the saved image.
"""
img_bytes = BytesIO(
to_sync(
generate_bytes_img(
"http://localhost:5000/index.html",
path="test_img.jpg", # Use of path to specify the type
selector="#godot",
)
)
)
img_bytes = Image.open(img_bytes)
img_file = Image.open("test_img.jpg")
self.assertTupleEqual(img_bytes.size, img_file.size)
self.assertAlmostEqual(_rmsd(img_bytes, img_file), 0, delta=0.05)
def test_screenshot_bad_url(self):
with self.assertRaises(BadUrl):
to_sync(generate_bytes_img("http://fake"))
def test_selector_manager(self):
"""
Tests _selector_manager
"""
url = "http://fake"
params_page1 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for": None,
"wait_for_xpath": None,
"wait_until": ["load"],
}
page1 = FakePage(
arg_viewport=None,
wait_until=["load"],
goto_called=True,
url=url,
waitForSelector_called=False,
wait_for=None,
)
new_page1 = to_sync(_selector_manager(page1, params_page1))
self.assertEqual(new_page1.arg_viewport, None)
self.assertEqual(new_page1.wait_until, ["load"])
self.assertTrue(new_page1.goto_called)
self.assertEqual(new_page1.url, url)
self.assertFalse(new_page1.waitForSelector_called)
self.assertEqual(new_page1.wait_for, None)
self.assertFalse(new_page1.querySelector_called)
self.assertEqual(new_page1.selector, None)
params_page2 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": "div",
"wait_for": None,
"wait_for_xpath": None,
"wait_until": ["load"],
}
page2 = FakePage(
arg_viewport=None,
wait_until=["load"],
goto_called=True,
url=url,
waitForSelector_called=False,
wait_for=None,
)
new_page2 = to_sync(_selector_manager(page2, params_page2))
self.assertEqual(new_page2.arg_viewport, None)
self.assertEqual(new_page2.wait_until, ["load"])
self.assertTrue(new_page2.goto_called)
self.assertEqual(new_page2.url, url)
self.assertFalse(new_page2.waitForSelector_called)
self.assertEqual(new_page2.wait_for, None)
self.assertTrue(new_page2.querySelector_called)
self.assertEqual(new_page2.selector, "div")
def main():
parser = ArgumentParser(description=__doc__)
# Mandatory arguments
parser.add_argument("url", help="url to screenshot")
parser.add_argument("path", help="path to image to be written")
# Optionnal arguments
parser.add_argument("-f",
"--fullpage",
action="store_true",
help="Take a screenshot of the whole scrollable page")
parser.add_argument("--width",
type=int,
metavar="width",
default=800,
help="Set the width of the page")
parser.add_argument("--height",
type=int,
metavar="height",
default=600,
help="Set the height of the page")
parser.add_argument(
"--wait_until",
nargs="+",
type=str,
choices=["load", "domcontentloaded", "networkidle0", "networkidle2"],
default="load",
help="How long do you want to wait for the page to be loaded")
# Credentials group
credentials_group = parser.add_argument_group(
title="Credentials (optional)")
credentials_group.add_argument('--username', help='The username to use')
credentials_group.add_argument('--password', help='The password to use')
credentials_group.add_argument('--token',
help='The header line to add. \
Must be like the following expression: key:token')
# CSS3 selectors group
selector_group = parser.add_argument_group(
title="CSS3 selectors (optional)",
description="Using quotes is recommended for \
complex CSS3 selectors")
selector_group.add_argument(
"--selector",
metavar="seletor",
help="The CSS3 selector of the element you want in the screenshot")
selector_group.add_argument(
"--wait_for",
metavar="wait_for_selector",
help="The CSS3 selector of an element you want to wait to be \
loaded before taking the screenshot")
# Browser group
browser_group = parser.add_argument_group(
title='Browser (optional)',
description='By default a headless browser is \
opened and closed')
browser_group.add_argument('--no-browser',
action='store_true',
default=False,
help='No browser is opened')
browser_group.add_argument('-g',
'--graphic',
dest="headless",
action="store_false",
default=True,
help="Open the browser in graphic mode")
browser_group.add_argument("--no-sandbox",
action="store_const",
const=["--no-sandbox"],
default="[]",
help="Open the browser without sandbox")
browser_group.add_argument(
'--no-close',
action='store_true',
default=False,
help='The browser is not closed when the job is done')
args = parser.parse_args()
credentials = None
if args.username and not args.password:
logger.error('A password must be specified')
exit(1)
elif not args.username and args.password:
logger.error('A username must be specified')
exit(1)
elif args.username and args.password:
credentials = {'username': args.username, 'password': args.password}
elif args.token:
regex_token = search(r'(?P<key>[^<]+)\:(?P<token>[^<]+)', args.token)
try:
key = regex_token.group('key')
token = regex_token.group('token')
credentials = {key: token, 'token_in_header': True}
except AttributeError as _:
logger.error('Bad token argument, please read the documentation')
exit(1)
if not args.no_browser:
to_sync(open_browser(args.headless, launch_args=args.no_sandbox))
to_sync(
generate_bytes_img(args.url,
path=args.path,
width=args.width,
height=args.height,
full_page=args.fullpage,
selector=args.selector,
wait_for=args.wait_for,
wait_until=args.wait_until,
credentials=credentials))
if not args.no_browser and not args.no_close:
to_sync(close_browser())
def tearDown(self):
to_sync(close_browser())
if exists("test_img.jpg"):
remove("test_img.jpg")
def test_page_manager(self, mock_get_token, mock_get_local_storage_token):
"""
Tests _page_manager
"""
browser = FakeBrowser()
url = "http://fake"
mock_get_token.return_value = {"token": "xxx"}
mock_get_local_storage_token.return_value = {"token": "xxx"}
params_page1 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for": None,
"wait_for_xpath": None,
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {},
"use_local_token": None,
}
page1 = to_sync(_page_manager(browser, url, params_page1))
self.assertEqual(page1.arg_viewport, None)
self.assertEqual(page1.credentials, None)
self.assertEqual(page1.wait_until, ["load"])
self.assertTrue(page1.goto_called)
self.assertEqual(page1.url, url)
self.assertFalse(page1.waitForSelector_called)
self.assertEqual(page1.wait_for, None)
self.assertFalse(page1.querySelector_called)
self.assertEqual(page1.selector, None)
params_page2 = {
"arg_viewport": {
"width": 800,
"height": 800
},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for": None,
"wait_for_xpath": None,
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {},
"use_local_token": None,
}
page2 = to_sync(_page_manager(browser, url, params_page2))
self.assertEqual(page2.arg_viewport, {"width": 800, "height": 800})
self.assertEqual(page2.credentials, None)
self.assertEqual(page2.wait_until, ["load"])
self.assertTrue(page2.goto_called)
self.assertEqual(page2.url, url)
self.assertFalse(page2.waitForSelector_called)
self.assertEqual(page2.wait_for, None)
self.assertFalse(page2.querySelector_called)
self.assertEqual(page2.selector, None)
params_page3 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for": None,
"wait_for_xpath": None,
"wait_until": ["load", "networkidle0"],
"credentials": {},
"credentials_token_request": {},
"use_local_token": None,
}
page3 = to_sync(_page_manager(browser, url, params_page3))
self.assertEqual(page3.arg_viewport, None)
self.assertEqual(page3.credentials, None)
self.assertEqual(page3.wait_until, ["load", "networkidle0"])
self.assertTrue(page3.goto_called)
self.assertEqual(page3.url, url)
self.assertFalse(page3.waitForSelector_called)
self.assertEqual(page3.wait_for, None)
self.assertFalse(page3.querySelector_called)
self.assertEqual(page3.selector, None)
params_page4 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for": "div",
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {},
"use_local_token": None,
}
page4 = to_sync(_page_manager(browser, url, params_page4))
self.assertEqual(page4.arg_viewport, None)
self.assertEqual(page4.credentials, None)
self.assertEqual(page4.wait_until, ["load"])
self.assertTrue(page4.goto_called)
self.assertEqual(page4.url, url)
self.assertTrue(page4.waitForSelector_called)
self.assertEqual(page4.wait_for, "div")
self.assertFalse(page4.querySelector_called)
self.assertEqual(page4.selector, None)
params_page5 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for": "div",
"wait_until": ["load"],
"credentials": {
"login": True,
"credentials_data": {
"username": "******",
"password": "******"
},
},
"credentials_token_request": {},
"use_local_token": None,
}
page5 = to_sync(_page_manager(browser, url, params_page5))
self.assertEqual(page5.arg_viewport, None)
self.assertEqual(page5.credentials, {
"username": "******",
"password": "******"
})
self.assertEqual(page5.wait_until, ["load"])
self.assertTrue(page5.goto_called)
self.assertEqual(page5.url, url)
self.assertTrue(page5.waitForSelector_called)
self.assertEqual(page5.wait_for, "div")
self.assertFalse(page5.querySelector_called)
self.assertEqual(page5.selector, None)
params_page6 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for": "div",
"wait_until": ["load"],
"credentials": {
"token_in_header": True,
"credentials_data": {
"token": "xxx"
},
},
"credentials_token_request": {},
"use_local_token": None,
}
page6 = to_sync(_page_manager(browser, url, params_page6))
self.assertEqual(page6.arg_viewport, None)
self.assertEqual(page6.credentials, {"token": "xxx"})
self.assertEqual(page6.wait_until, ["load"])
self.assertTrue(page6.goto_called)
self.assertEqual(page6.url, url)
self.assertTrue(page6.waitForSelector_called)
self.assertEqual(page6.wait_for, "div")
self.assertFalse(page6.querySelector_called)
self.assertEqual(page6.selector, None)
params_page7 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for_xpath": "div",
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {},
"use_local_token": None,
}
page7 = to_sync(_page_manager(browser, url, params_page7))
self.assertEqual(page7.arg_viewport, None)
self.assertEqual(page7.credentials, None)
self.assertEqual(page7.wait_until, ["load"])
self.assertTrue(page7.goto_called)
self.assertEqual(page7.url, url)
self.assertFalse(page7.waitForSelector_called)
self.assertEqual(page7.wait_for, None)
self.assertTrue(page7.waitForxPath_called)
self.assertEqual(page7.wait_for_xpath, "div")
self.assertFalse(page7.querySelector_called)
self.assertEqual(page7.selector, None)
params_page8 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for_xpath": "div",
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {
"url": "http://fake",
"username": "******",
"password": "******",
"local_storage": False
},
"use_local_token": None,
}
page8 = to_sync(_page_manager(browser, url, params_page8))
self.assertEqual(page8.arg_viewport, None)
self.assertEqual(page8.credentials, {"token": "xxx"})
self.assertEqual(page8.wait_until, ["load"])
self.assertTrue(page8.goto_called)
self.assertEqual(page8.url, url)
self.assertFalse(page8.waitForSelector_called)
self.assertEqual(page8.wait_for, None)
self.assertTrue(page8.waitForxPath_called)
self.assertEqual(page8.wait_for_xpath, "div")
self.assertFalse(page8.querySelector_called)
self.assertEqual(page8.selector, None)
params_page9 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for_xpath": "div",
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {
"url": "http://fake",
"username": "******",
"password": "******",
"local_storage": True
},
"use_local_token": True,
}
page9 = to_sync(_page_manager(browser, url, params_page9))
self.assertEqual(page9.arg_viewport, None)
self.assertEqual(page9.credentials, {"token": "xxx"})
self.assertEqual(page9.wait_until, ["load"])
self.assertTrue(page9.goto_called)
self.assertEqual(page9.url, url)
self.assertFalse(page9.waitForSelector_called)
self.assertEqual(page9.wait_for, None)
self.assertTrue(page9.waitForxPath_called)
self.assertEqual(page9.wait_for_xpath, "div")
self.assertFalse(page9.querySelector_called)
self.assertEqual(page9.selector, None)
params_page10 = {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for_xpath": "div",
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {},
"use_local_token": True,
}
page10 = to_sync(
_page_manager(
browser, url, {
"arg_viewport": {},
"screenshot_options": {
"fullPage": False
},
"selector": None,
"wait_for_xpath": "div",
"wait_until": ["load"],
"credentials": {},
"credentials_token_request": {
"url": "http://fake",
"username": "******",
"password": "******",
"local_storage": True
},
"use_local_token": True,
}))
page10 = to_sync(_page_manager(browser, url, params_page10))
self.assertEqual(page10.arg_viewport, None)
self.assertEqual(page10.credentials, {"token": "xxx"})
self.assertEqual(page10.wait_until, ["load"])
self.assertTrue(page10.goto_called)
self.assertEqual(page10.url, url)
self.assertFalse(page10.waitForSelector_called)
self.assertEqual(page10.wait_for, None)
self.assertTrue(page10.waitForxPath_called)
self.assertEqual(page10.wait_for_xpath, "div")
self.assertFalse(page10.querySelector_called)
self.assertEqual(page10.selector, None)
