diff --git a/util/topgen/entropy_buffer_generator.py b/util/topgen/entropy_buffer_generator.py
index 723114a..934b5e6 100755
--- a/util/topgen/entropy_buffer_generator.py
+++ b/util/topgen/entropy_buffer_generator.py
@@ -50,13 +50,10 @@
     return parser.parse_args()
 
 
-def main():
-
-    args = parse_args()
-    k = args.num_bytes
-    out = args.output_file
-    sec = args.sec
-    seed = args.seed
+def gen_buffer(k: int,
+               out,
+               sec: bool,
+               seed: int):
 
     if (sec and seed):
         log.error("Options --sec and --seed cannot be used together")
@@ -80,5 +77,16 @@
             fp.write("%s\n" % item)
 
 
+def main():
+
+    args = parse_args()
+    k = args.num_bytes
+    out = args.output_file
+    sec = args.sec
+    seed = args.seed
+
+    gen_buffer(k, out, sec, seed)
+
+
 if __name__ == "__main__":
     main()