diff --git a/tests/tests_oss_simple_flow.py b/tests/tests_oss_simple_flow.py
index 8088456c..f6a7c0c5 100644
--- a/tests/tests_oss_simple_flow.py
+++ b/tests/tests_oss_simple_flow.py
@@ -4,8 +4,7 @@
 from mbdirector.benchmark import Benchmark
 from mbdirector.runner import RunConfig
 
-import logging
-    
+
 def test_preload_and_set_get(env):
     key_max = 500000
     key_min = 1
@@ -368,14 +367,14 @@ def test_default_arbitrary_command_hset_multi_data_placeholders(env):
     benchmark_specs["args"].append('--command=HSET __key__ field1 __data__ field2 __data__ field3 __data__')
     config = get_default_memtier_config()
     master_nodes_list = env.getMasterNodesList()
-    overall_expected_request_count = get_expected_request_count(memtier_config)
+    overall_expected_request_count = get_expected_request_count(config)
 
-    add_required_env_arguments(benchmark_specs, memtier_config, env, master_nodes_list)
+    add_required_env_arguments(benchmark_specs, config, env, master_nodes_list)
 
     # Create a temporary directory
     test_dir = tempfile.mkdtemp()
 
-    config = RunConfig(test_dir, env.testName, memtier_config, {})
+    config = RunConfig(test_dir, env.testName, config, {})
     ensure_clean_benchmark_folder(config.results_dir)
 
     benchmark = Benchmark.from_json(config, benchmark_specs)
@@ -430,12 +429,13 @@ def test_default_set_get_rate_limited(env):
             overall_request_count = agg_info_commandstats(master_nodes_connections, merged_command_stats)
             assert_minimum_memtier_outcomes(config, env, memtier_ok, overall_expected_request_count, overall_request_count, request_delta)
 
+
 def test_data_import(env):
     env.skipOnCluster()
     benchmark_specs = {"name": env.testName, "args": [f"--data-import={ROOT_FOLDER}/tests/data-import-2-keys.txt",'--ratio=1:1']}
     addTLSArgs(benchmark_specs, env)
     config = get_default_memtier_config()
-    master_nodes_list = env.getMasterNodesList()  
+    master_nodes_list = env.getMasterNodesList()
     overall_expected_request_count = get_expected_request_count(config,1, 2)
     add_required_env_arguments(benchmark_specs, config, env, master_nodes_list)
 
@@ -456,7 +456,7 @@ def test_data_import(env):
     memtier_ok = benchmark.run()
 
     assert_keyspace_range(env, 2, 1, master_nodes_connections)
-    
+
     merged_command_stats = {'cmdstat_set': {'calls': 0}, 'cmdstat_get': {'calls': 0}}
     overall_request_count = agg_info_commandstats(master_nodes_connections, merged_command_stats)
     assert_minimum_memtier_outcomes(config, env, memtier_ok, overall_expected_request_count, overall_request_count)
@@ -467,7 +467,7 @@ def test_data_import_setex(env):
     benchmark_specs = {"name": env.testName, "args": [f"--data-import={ROOT_FOLDER}/tests/data-import-2-keys-expiration.txt",'--ratio=1:1']}
     addTLSArgs(benchmark_specs, env)
     config = get_default_memtier_config()
-    master_nodes_list = env.getMasterNodesList()  
+    master_nodes_list = env.getMasterNodesList()
     overall_expected_request_count = get_expected_request_count(config,1, 2)
     add_required_env_arguments(benchmark_specs, config, env, master_nodes_list)
 
@@ -488,8 +488,7 @@ def test_data_import_setex(env):
     memtier_ok = benchmark.run()
 
     assert_keyspace_range(env, 2, 1, master_nodes_connections)
-    
+
     merged_command_stats = {'cmdstat_setex': {'calls': 0}, 'cmdstat_get': {'calls': 0}}
     overall_request_count = agg_info_commandstats(master_nodes_connections, merged_command_stats)
     assert_minimum_memtier_outcomes(config, env, memtier_ok, overall_expected_request_count, overall_request_count)
-