diff --git a/data_generator.py b/data_generator.py
index b6373b1..096da3b 100644
--- a/data_generator.py
+++ b/data_generator.py
@@ -1,11 +1,11 @@
-import sys
 import os
+import sys
 import uuid
 import random
 import itertools
 
 from petname import Generate as pet_generate
-from storage_api import create_storage, TEST_PATH
+from storage_api import create_storage
 
 from report import ssize_to_kb
 
@@ -13,12 +13,11 @@
 random.shuffle(types)
 tp = itertools.cycle(types)
 
-sz = ["4k", "64k", "1m"]
+sz = ["1k", "4k", "64k", "256k", "1m"]
 op_type = ["randread", "read", "randwrite", "write"]
 is_sync = ["s", "a"]
 
-
-storage = create_storage("file://" + TEST_PATH + "/sample.json", "", "")
+storage = create_storage(sys.argv[1], "", "")
 combinations = list(itertools.product(op_type, is_sync, sz))
 
 for i in range(30):