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)