<reponame>vrserpa/pandas-dev-flaker
import ast
import tokenize
from io import StringIO

import pytest

from pandas_dev_flaker.__main__ import run


def results(s):
    return {
        "{}:{}: {}".format(*r)
        for r in run(
            ast.parse(s),
            list(tokenize.generate_tokens(StringIO(s).readline)),
        )
    }


@pytest.mark.parametrize(
    "source",
    (
        pytest.param(
            "from foo import remove\nremove('bar')",
            id="import remove from foo",
        ),
        pytest.param(
            "import foo\nfoo.remove('foo')",
            id="foo.remove",
        ),
    ),
)
def test_noop(source):
    assert not results(source)


@pytest.mark.parametrize(
    "source, expected",
    (
        pytest.param(
            "import os\nos.remove('foo')",
            "2:0: PDF008 found 'os.remove' (use 'tm.ensure_clean' instead)",
            id="os.remove call",
        ),
        pytest.param(
            "from os import remove\nremove('foo')",
            "2:0: PDF008 found 'os.remove' (use 'tm.ensure_clean' instead)",
            id="import remove from os",
        ),
    ),
)
def test_violation(source, expected):
    (result,) = results(source)
    assert result == expected
