added corresponding tests to the new method of the multiple subscription

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