zeek/auxil/broker/tests/python/forwarding.py
Patrick Kelley 8fd444092b initial
2025-05-07 15:35:15 -04:00

94 lines
2.5 KiB
Python

import unittest
import broker
def cleanup(es, ss):
for s in ss:
if s:
s.reset()
for e in es:
e.shutdown()
def setup_peers(
opts1=None,
opts2=None,
opts3=None,
opts4=None,
create_s1=True,
create_s2=True,
create_s3=True,
create_s4=True,
):
def cfg(opts):
return (
broker.Configuration(opts)
if opts
else broker.Configuration(broker.BrokerOptions())
)
ep1 = broker.Endpoint(cfg(opts1))
ep2 = broker.Endpoint(cfg(opts2))
ep3 = broker.Endpoint(cfg(opts3))
ep4 = broker.Endpoint(cfg(opts4))
s1 = ep1.make_subscriber("/test/") if create_s1 else None
s2 = ep2.make_subscriber("/test/") if create_s2 else None
s3 = ep3.make_subscriber("/test/") if create_s3 else None
s4 = ep4.make_subscriber("/test/") if create_s4 else None
p2 = ep2.listen("127.0.0.1", 0)
p3 = ep3.listen("127.0.0.1", 0)
p4 = ep4.listen("127.0.0.1", 0)
# ep1 <-> ep2 <-> ep3 <-> ep4
ep1.peer("127.0.0.1", p2, 1.0)
ep2.peer("127.0.0.1", p3, 1.0)
ep3.peer("127.0.0.1", p4, 1.0)
return ((ep1, ep2, ep3, ep4), (s1, s2, s3, s4))
class TestCommunication(unittest.TestCase):
def test_two_subscribed_hops(self):
# Two hops that are subscribed.
((ep1, ep2, ep3, ep4), (s1, s2, s3, s4)) = setup_peers()
ep1.await_peer(ep4.node_id())
ep4.await_peer(ep1.node_id())
ep1.publish("/test/foo", "Foo!")
ep4.publish("/test/bar", "Bar!")
x = s4.get()
self.assertEqual(x, ("/test/foo", "Foo!"))
x = s1.get()
self.assertEqual(x, ("/test/bar", "Bar!"))
cleanup((ep1, ep2, ep3, ep4), (s1, s2, s3, s4))
#### Note: disabled until we switch back to source-routing.
#
# def test_two_unsubscribed_hops(self):
# # Two hops that are not subscribed, but still forward due to the source
# # routing.
# ((ep1, ep2, ep3, ep4), (s1, s2, s3, s4)) = setup_peers(create_s2=False, create_s3=False)
#
# ep1.await_peer(ep4.node_id())
# ep4.await_peer(ep1.node_id())
#
# ep1.publish("/test/foo", "Foo!")
# ep4.publish("/test/bar", "Bar!")
#
# x = s4.get()
# self.assertEqual(x, ('/test/foo', 'Foo!'))
# x = s1.get()
# self.assertEqual(x, ('/test/bar', 'Bar!'))
# cleanup((ep1, ep2, ep3, ep4), (s1, s2, s3, s4))
if __name__ == "__main__":
# TestCommunication().test_two_hops()
unittest.main(verbosity=3)