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_heuristics(self):
a = HeuristicA()
b = HeuristicB()
h = Heuristics()
h.add(a)
h.add(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)
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)
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)
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)