# tests/test_load_test.py
import pytest
from unittest.mock import Mock, patch
from multiprocessing import Queue
from llm_load_test.load_test import run_main_process, gather_results, stop_procs, create_procs

def test_run_main_process():
    # Mock dependencies
    dataset = Mock()
    dataset.get_next_n_queries.return_value = ["query1", "query2"]
    dataset_q = Queue()
    stop_q = Queue()
    
    # Test normal execution
    with patch('time.time') as mock_time, patch('time.sleep'):
        mock_time.side_effect = [0, 0.5, 1, 1.5, 2]  # Simulate time progression
        run_main_process(2, 1, dataset, dataset_q, stop_q)
        
        # Verify dataset queue was filled
        assert dataset_q.qsize() == 0  # Queue should be empty after cleanup
        assert stop_q.qsize() == 1  # Stop signal should be sent

def test_gather_results():
    # Mock pipes
    pipe1, pipe2 = Mock(), Mock()
    pipe1.recv.return_value = ["result1", "result2"]
    pipe2.recv.return_value = ["result3"]
    
    results = gather_results([pipe1, pipe2])
    assert len(results) == 3
    assert "result1" in results
    assert "result3" in results

def test_stop_procs():
    # Mock processes and queues
    proc1, proc2 = Mock(), Mock()
    dataset_q = Queue()
    stop_q = Queue()
    
    # Add some items to queues
    dataset_q.put("item1")
    
    stop_procs([proc1, proc2], dataset_q, stop_q)
    
    # Verify cleanup
    assert proc1.join.called
    assert proc2.join.called
    assert dataset_q.empty()
    assert stop_q.empty()

# tests/test_plugins.py
import pytest
from llm_load_test.plugins.dummy_plugin import DummyPlugin
from llm_load_test.result import RequestResult

def test_dummy_plugin_request():
    # Test non-streaming request
    plugin = DummyPlugin({"streaming": False})
    query = {
        "text": "hello world",
        "output_tokens": 5,
        "input_tokens": 2
    }
    
    result = plugin.request_http(query, user_id=1)
    assert isinstance(result, RequestResult)
    assert result.output_text == "dlrow olleh"  # Reversed text
    assert result.output_tokens == 5
    
def test_dummy_plugin_streaming():
    # Test streaming request
    plugin = DummyPlugin({"streaming": True})
    query = {
        "text": "hello world",
        "output_tokens": 5,
        "input_tokens": 2
    }
    
    result = plugin.streaming_request_http(query, user_id=1)
    assert isinstance(result, RequestResult)
    assert result.ack_time is not None
    assert result.first_token_time is not None
    assert result.end_time > result.start_time

# tests/test_result.py
import pytest
from llm_load_test.result import RequestResult

def test_request_result_calculation():
    result = RequestResult(user_id=1, input_id="test1", input_tokens=10)
    result.start_time = 1000
    result.ack_time = 1001
    result.first_token_time = 1002
    result.end_time = 1005
    result.output_tokens = 20
    
    result.calculate_results()
    
    assert result.time_to_first_token == 2  # 1002 - 1000
    assert result.total_time == 5  # 1005 - 1000
    assert result.tokens_per_second == 4  # 20 tokens / 5 seconds

def test_request_result_error_handling():
    result = RequestResult(user_id=1, input_id="test1", input_tokens=10)
    result.error_code = 500
    result.error_text = "Internal Server Error"
    
    assert result.has_error()
    assert result.error_code == 500

# tests/test_integration.py
import pytest
from llm_load_test.load_test import main
from llm_load_test.plugins.dummy_plugin import DummyPlugin

@pytest.fixture
def mock_config(tmp_path):
    config_file = tmp_path / "test_config.yaml"
    config_file.write_text("""
plugin: "dummy"
plugin_options:
  streaming: false
dataset:
  type: "simple"
  input_text: "test query"
  input_tokens: 2
  output_tokens: 5
load_options:
  concurrency: 2
  duration: 1
""")
    return str(config_file)

def test_complete_flow(mock_config):
    with pytest.raises(SystemExit) as exc_info:
        main(["-c", mock_config, "--log_level", "debug"])
    assert exc_info.value.code == 0