diff --git a/tests/core/test_subscribe.py b/tests/core/test_subscribe.py index f3002b7..e95ba16 100644 --- a/tests/core/test_subscribe.py +++ b/tests/core/test_subscribe.py @@ -6,6 +6,8 @@ from kube_hunter.core.events.types import Event, Service from kube_hunter.core.events import handler counter = 0 +first_run = True + set_config(Config()) @@ -13,18 +15,22 @@ class OnceOnlyEvent(Service, Event): def __init__(self): Service.__init__(self, "Test Once Service") + class RegularEvent(Service, Event): def __init__(self): Service.__init__(self, "Test Service") + class AnotherRegularEvent(Service, Event): def __init__(self): Service.__init__(self, "Test Service (another)") + class DifferentRegularEvent(Service, Event): def __init__(self): Service.__init__(self, "Test Service (different)") + @handler.subscribe_once(OnceOnlyEvent) class OnceHunter(Hunter): def __init__(self, event): @@ -38,16 +44,38 @@ class RegularHunter(Hunter): global counter counter += 1 + @handler.subscribe_many([DifferentRegularEvent, AnotherRegularEvent]) class SmartHunter(Hunter): + def __init__(self, events): + global counter, first_run + counter += 1 + + # we add an attribute on the second scan. + # here we test that we get the latest event + different_event = events.get_by_class(DifferentRegularEvent) + if first_run: + first_run = False + assert different_event.new_value == None + else: + assert different_event.new_value + + +@handler.subscribe_many([DifferentRegularEvent, AnotherRegularEvent]) +class SmartHunter2(Hunter): def __init__(self, events): global counter counter += 1 + # check if we can access the events + assert events.get_by_class(DifferentRegularEvent).__class__ == DifferentRegularEvent + assert events.get_by_class(AnotherRegularEvent).__class__ == AnotherRegularEvent + + def test_subscribe_mechanism(): global counter counter = 0 - + # first test normal subscribe and publish works handler.publish_event(RegularEvent()) handler.publish_event(RegularEvent()) @@ -56,6 +84,7 @@ def test_subscribe_mechanism(): time.sleep(0.02) assert counter == 3 + def test_subscribe_once_mechanism(): global counter counter = 0 @@ -81,23 +110,20 @@ def test_subscribe_many_mechanism(): # testing the multiple subscription mechanism handler.publish_event(DifferentRegularEvent()) - handler.publish_event(AnotherRegularEvent()) handler.publish_event(DifferentRegularEvent()) + handler.publish_event(DifferentRegularEvent()) + handler.publish_event(DifferentRegularEvent()) + handler.publish_event(DifferentRegularEvent()) + handler.publish_event(AnotherRegularEvent()) time.sleep(0.02) - # We expect that SmartHunter to run once and RegularEvent to run once. + # We expect SmartHunter and SmartHunter2 to be executed once. hence the counter should be 2 assert counter == 2 counter = 0 - handler.publish_event(AnotherRegularEvent()) - handler.publish_event(AnotherRegularEvent()) - handler.publish_event(AnotherRegularEvent()) - handler.publish_event(DifferentRegularEvent()) - handler.publish_event(DifferentRegularEvent()) - handler.publish_event(DifferentRegularEvent()) + # Test using most recent event + newer_version_event = DifferentRegularEvent() + newer_version_event.new_value = True + handler.publish_event(newer_version_event) - time.sleep(0.02) - # (Regular, Another) or (Another, Regular) sequences trigger the SmartHunter. - # Regular trigger the RegularHunter. - # OnceHunter should not be triggered here. - assert counter == 1 + assert counter == 2