fix errors in tests

Author: XuehengWang

tests/test_cli.py

@@ -81,7 +81,7 @@ def test_cli_invalid_pattern(self):
         )
 
         self.assertNotEqual(result.exit_code, 0)
-        self.assertIn("Invalid pattern", result.output)
+        self.assertIn("Invalid value for '--pattern'", result.output)
 
 
 if __name__ == "__main__":

tests/test_data_loader.py

@@ -28,7 +28,7 @@ def test_remote_files(self):
         2. loading csv format with pandas
         """
         datasets, sort = load_datasets(
-            "../tracestorm/datasets_config/datasets_config_hf.json"
+            "tracestorm/datasets_config/datasets_config_hf.json"
         )
         assert isinstance(datasets, list)
         assert isinstance(datasets[0], Dataset) and isinstance(
@@ -42,7 +42,7 @@ def test_remote_files(self):
     def test_missing_fields(self):
         """Test loading with missing sort strategy and selection ratio in the config file"""
         datasets, sort = load_datasets(
-            "../tracestorm/datasets_config/datasets_config_missing.json"
+            "tracestorm/datasets_config/datasets_config_missing.json"
         )
         assert isinstance(datasets, list) and len(datasets) == 2
         assert sort == "random"

tracestorm/cli.py

@@ -1,5 +1,5 @@
 import os
-from typing import Tuple
+from typing import Tuple, Optional
 
 import click
 
@@ -21,7 +21,7 @@
 
 
 def create_trace_generator(
-    pattern: str, rps: int, duration: int, seed: int
+    pattern: str, rps: int, duration: int, seed: Optional[int] = None
 ) -> Tuple[TraceGenerator, str]:
     """
     Create appropriate trace generator based on pattern and validate parameters.

tracestorm/datasets_config/datasets_config_missing.json

@@ -5,6 +5,6 @@
     },
     "dataset_2": {
         "file_name": "lmsys/chatbot_arena_conversations",
-        "prompt_field": "user_prompt"
+        "prompt_field": "conversation_a"
     }
 }