def test_assign_engine_and_kb_on_heuristics(self):
a = HeuristicA()
b = HeuristicB()
h = Heuristics(request_engine="X", kb="Y")
h.add_multiple([a, b])
self.assertEqual(a.engine_set, "X")
self.assertEqual(b.kb_set, "Y")
def test_do_not_override_with_none(self):
a = HeuristicA()
b = HeuristicB()
h = Heuristics()
h.add_multiple([a, b])
self.assertEqual(a.engine_set, "NO")
self.assertEqual(b.kb_set, "NO")
def test_init_child_heuristics_on_add(self):
h = Heuristics(kb="kb", request_engine="engine")
parent_heuristic = ParentHeuristic()
h.add(parent_heuristic)
self.assertEqual(parent_heuristic.child_heuristics.kb, "kb")
self.assertEqual(parent_heuristic.child_heuristics.request_engine,
"engine")
def test_load_multiple_heuristics(self):
a = HeuristicA()
b = HeuristicB()
h = Heuristics()
h.add_multiple([a, b])
self.assertEqual(len(h.rulesets["before_request"]), 2)
self.assertEqual(len(h.rulesets["after_headers"]), 1)
self.assertEqual(len(h.rulesets["after_response"]), 1)
self.assertEqual(h.rulesets["before_request"].accept, h.before_request)
class Pipeline:
def __init__(self, *, engine=None):
self.kb = KnowledgeBase()
self.heuristics = Heuristics(kb=self.kb, request_engine=engine)
self.child_heuristics = Heuristics(kb=self.kb, request_engine=engine)
def add(self, heuristic, *, with_child=False):
self.heuristics.add(heuristic)
if with_child:
self.child_heuristics.add(heuristic)
async def perform_ok(self, entry):
await self.heuristics.before_request(entry)
await self.heuristics.after_headers(entry)
await self.heuristics.after_response(entry)
await self.heuristics.on_request_successful(entry)
async def test_default_ssl_parameters(self, loop):
engine = AioHttpEngine(loop=loop)
engine.session.request = make_mocked_coro(return_value=FakeResponse())
entry = Entry.create("http://www.example.com/")
await engine.perform(entry, Heuristics())
engine.session.request.assert_called_once_with(
method=entry.request.method,
url=entry.request.url,
allow_redirects=False,
timeout=0.2,
ssl=None)
async def test_perform_use_proxy_for_request(self, loop):
asyncio.set_event_loop(loop)
engine = AioHttpEngine(loop=loop, proxy="http://some.proxy.com/")
engine.session.request = make_mocked_coro(return_value=FakeResponse())
entry = Entry.create("http://www.example.com/1")
await engine.perform(entry, Heuristics())
engine.session.request.assert_called_once_with(
method=entry.request.method,
url="http://www.example.com/1",
timeout=0.2,
proxy="http://some.proxy.com/")
async def test_specify_header(self, loop):
asyncio.set_event_loop(loop)
engine = AioHttpEngine(loop=loop)
engine.session.request = make_mocked_coro(return_value=FakeResponse())
entry = Entry.create("http://www.example.com/1",
headers={"User-Agent": "Hammertime 1.2.3"})
await engine.perform(entry, Heuristics())
engine.session.request.assert_called_once_with(
method=entry.request.method,
url="http://www.example.com/1",
timeout=ANY,
allow_redirects=False,
headers={"User-Agent": "Hammertime 1.2.3"})
async def test_perform_use_timeout_of_entry_if_not_none(self, loop):
asyncio.set_event_loop(loop)
engine = AioHttpEngine(loop=loop, proxy="http://some.proxy.com/")
engine.session.request = make_mocked_coro(return_value=FakeResponse())
entry = Entry.create("http://www.example.com/1",
arguments={"timeout": 10})
await engine.perform(entry, Heuristics())
engine.session.request.assert_called_once_with(
method=entry.request.method,
url="http://www.example.com/1",
timeout=10,
proxy="http://some.proxy.com/",
allow_redirects=False)
def test_add_load_kb_if_value_already_defined_in_kb(self):
kb = KnowledgeBase()
h = Heuristics(kb=kb, request_engine="engine")
parent_heuristic = ParentHeuristic()
h.add(parent_heuristic)
h.add(HeuristicB())
child_heuristic = HeuristicB()
parent_heuristic.child_heuristics.add(child_heuristic)
self.assertEqual(child_heuristic.kb_load, kb)
async def test_set_ssl_parameters_set_certification_authority_certificate(
self, loop):
ssl_context = "ssl_context"
with patch(
"ssl.create_default_context",
MagicMock(return_value=ssl_context)) as create_default_context:
engine = AioHttpEngine(loop=loop)
engine.session.request = make_mocked_coro(
return_value=FakeResponse())
entry = Entry.create("http://www.example.com/")
engine.set_ssl_parameters(ca_certificate_file="certificate.cer")
await engine.perform(entry, Heuristics())
engine.session.request.assert_called_once_with(
method=entry.request.method,
url=entry.request.url,
allow_redirects=False,
timeout=0.2,
ssl=ssl_context)
create_default_context.assert_called_once_with(
cafile="certificate.cer")
def test_load_bad_heuristic(self):
a = HeuristicBad()
h = Heuristics()
with self.assertRaises(ValueError):
h.add(a)
def test_load_not_async(self):
a = HeuristicNonAsync()
h = Heuristics()
with self.assertRaises(ValueError):
h.add(a)
def __init__(self, *, max_redirect=15):
self.max_redirect = max_redirect
self.engine = None
self.child_heuristics = Heuristics()
def __init__(self, *, engine=None):
self.kb = KnowledgeBase()
self.heuristics = Heuristics(kb=self.kb, request_engine=engine)
self.child_heuristics = Heuristics(kb=self.kb, request_engine=engine)