from core.languages.python.testgen.preprocessing.preprocess_pipeline import preprocessing_testgen_pipeline
from core.languages.python.testgen.preprocessing.torch_tensor_limit import (
    OBJECT_MEMORY_LIMIT_MB,
    detect_torch_usage,
    get_tensor_size_note,
)


def test_detect_torch_usage_with_torch():
    """Test detect_torch_usage returns True when 'torch' is in the test context."""
    test_context = "import torch\nx = torch.rand(5, 5)"
    result = detect_torch_usage(test_context)
    assert result is True


def test_detect_torch_usage_with_tensor():
    """Test detect_torch_usage returns True when 'Tensor' is in the test context."""
    test_context = "from torch import Tensor\nx = Tensor(5, 5)"
    result = detect_torch_usage(test_context)
    assert result is True


def test_detect_torch_usage_without_torch():
    """Test detect_torch_usage returns False when neither 'torch' nor 'Tensor' is in the test context."""
    test_context = "import numpy as np\nx = np.array([1, 2, 3])"
    result = detect_torch_usage(test_context)
    assert result is False


def test_get_tensor_size_note_with_torch():
    """Test get_tensor_size_note returns a note when torch is detected."""
    test_context = "import torch\nx = torch.rand(5, 5)"
    result = get_tensor_size_note(test_context)
    expected = [f"Do not generate torch tensors with size greater than {OBJECT_MEMORY_LIMIT_MB}MB."]
    assert result == expected


def test_get_tensor_size_note_without_torch():
    """Test get_tensor_size_note returns empty list when torch is not detected."""
    test_context = "import numpy as np\nx = np.array([1, 2, 3])"
    result = get_tensor_size_note(test_context)
    assert not result


def test_get_all_notes_with_torch():
    """Test get_all_notes returns a list with a note when torch is detected."""
    test_context = "import torch\nx = torch.rand(5, 5)"
    result = get_tensor_size_note(test_context)
    expected = [f"Do not generate torch tensors with size greater than {OBJECT_MEMORY_LIMIT_MB}MB."]
    assert result == expected


def test_get_all_notes_without_torch():
    """Test get_all_notes returns an empty list when torch is not detected."""
    test_context = "import numpy as np\nx = np.array([1, 2, 3])"
    result = get_tensor_size_note(test_context)
    assert result == []


def test_preprocessing_pipeline():
    """Test the preprocessing pipeline with and without torch detection."""
    # Test with torch
    function_code = "import torch\ndef my_func():\n    return torch.rand(5, 5)"
    result = preprocessing_testgen_pipeline(function_code)
    expected = [f"Do not generate torch tensors with size greater than {OBJECT_MEMORY_LIMIT_MB}MB."]
    assert result == expected

    # Test without torch
    function_code = "import numpy as np\ndef my_func():\n    return np.array([1, 2, 3])"
    result = preprocessing_testgen_pipeline(function_code)
    expected = []
    assert result == expected
