from inspect_evals.humaneval.humaneval import extract_function_body, find_code


class TestFindCode:
    def test_code_block_with_language(self):
        completion = """```python
def func():
    return 1
```
"""
        assert find_code(completion) == "    return 1\n"

    def test_code_block_without_language(self):
        completion = """```
def func():
    return 1
```
"""
        assert find_code(completion) == "    return 1\n"

    def test_multiple_code_blocks(self):
        completion = """Here's one function:
```python
def func1():
    return 1
```
And here's another:
```python
def func2():
    return 2
```
"""
        assert find_code(completion) == "    return 1\n"

    def test_no_code_blocks(self):
        completion = "This is just plain text without any code blocks."
        assert find_code(completion) == completion

    def test_empty_code_block(self):
        completion = """```python
```"""
        assert find_code(completion) == ""


class TestExtractFunctionBody:
    def test_extract_simple_function(self):
        code = """
def simple_function():
    return "Hello, World!"
"""
        result = extract_function_body(code)
        assert result == '    return "Hello, World!"\n'

    def test_extract_function_with_arguments(self):
        code = """
def function_with_args(a, b=10):
    result = a + b
    return result
"""
        result = extract_function_body(code)
        assert result == "    result = a + b\n    return result\n"

    def test_function_with_comments(self):
        code = """
def function_with_comments():
    # This is a comment
    x = 10  # Inline comment
    return x
"""
        result = extract_function_body(code)
        lines = result.split("\n")
        assert lines[0] == "    # This is a comment"
        assert lines[1] == "    x = 10  # Inline comment"
        assert lines[2] == "    return x"
