Commit b55d1833 authored by Jarrod Pas's avatar Jarrod Pas
Browse files

Support dict in trace and traffic

parent 19fc7757
......@@ -469,6 +469,9 @@ def random_trace(nodes, seed=None, step=1):
"""Generate a random contact trace."""
random = Random(seed)
if isinstance(nodes, dict):
nodes = list(nodes)
for now in count(step=step):
join = bool(random.getrandbits(1))
node_a, node_b = random.sample(nodes, 2)
......@@ -494,13 +497,16 @@ Traffic = namedtuple('Traffic', [
])
def random_traffic(nodes, start=0, step=1, seed=None, **options):
def random_traffic(nodes, start=0, speed=1, seed=None, **options):
"""Generate traffic from random source to random destination every step."""
random = Random(seed)
if isinstance(nodes, dict):
nodes = list(nodes)
time_to_live = options.get('time_to_live', float('inf'))
payload = options.get('payload', 0)
for created in count(start=start, step=step):
for created in count(start=start, step=speed):
source, destination = random.sample(nodes, 2)
yield Traffic(source, destination, created, time_to_live, payload)
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment