traffic_replay: Add a max-members option to cap group size

[samba.git] / script / traffic_replay

diff --git a/script/traffic_replay b/script/traffic_replay
index 617a746be1bdb9b501fff6794f5dc48f9bb55045..0ee0f9b65752291a38ecb256acb28e2217fda71a 100755 (executable)
--- a/script/traffic_replay
+++ b/script/traffic_replay
@@ -112,6 +112,8 @@ def main():
     user_gen_group.add_option('--group-memberships', type='int', default=0,
                               help='Total memberships to assign across all '
                               'test users and all groups')
+    user_gen_group.add_option('--max-members', type='int', default=None,
+                              help='Max users to add to any one group')
     parser.add_option_group(user_gen_group)
 
     sambaopts = options.SambaOptions(parser)
@@ -254,7 +256,7 @@ def main():
                 duration = opts.duration
 
         except ValueError as e:
-            if not e.message.startswith('need more than'):
+            if not str(e).startswith('need more than'):
                 raise
 
             model = traffic.TrafficModel()
@@ -324,12 +326,18 @@ def main():
         sys.exit(1)
 
     if opts.generate_users_only:
+        # generate computer accounts for added realism. Assume there will be
+        # some overhang with more computer accounts than users
+        computer_accounts = int(1.25 * number_of_users)
         traffic.generate_users_and_groups(ldb,
                                           opts.instance_id,
                                           opts.fixed_password,
                                           opts.number_of_users,
                                           opts.number_of_groups,
-                                          opts.group_memberships)
+                                          opts.group_memberships,
+                                          opts.max_members,
+                                          machine_accounts=computer_accounts,
+                                          traffic_accounts=False)
         sys.exit()
 
     tempdir = tempfile.mkdtemp(prefix="samba_tg_")
@@ -340,7 +348,10 @@ def main():
                                       opts.fixed_password,
                                       number_of_users,
                                       opts.number_of_groups,
-                                      opts.group_memberships)
+                                      opts.group_memberships,
+                                      opts.max_members,
+                                      machine_accounts=len(conversations),
+                                      traffic_accounts=True)
 
     accounts = traffic.generate_replay_accounts(ldb,
                                                 opts.instance_id,