import pytest

from bluesky.utils import register_transform


@pytest.fixture
def transform_cell():
    IPython = pytest.importorskip("IPython")
    ip = IPython.core.interactiveshell.InteractiveShell()
    register_transform("RE", prefix="<", ip=ip)
    if IPython.__version__ >= "7":
        return ip.transform_cell
    else:
        return ip.input_splitter.transform_cell


def test_register_transform_smoke(transform_cell):
    assert True


@pytest.mark.parametrize("cell", ["a < b\n", "RE(plan(a < b))\n", "for j in range(5):\n    < a\n    <b\n"])
def test_no_transform(transform_cell, cell):
    new_cell = transform_cell(cell)
    assert cell == new_cell


@pytest.mark.parametrize("cell", ["< b\n", "<a\n", '<fly(a, b, c), md={"a": "aardvark"}\n', "  <b\n", "<a, c=d\n"])
def test_transform(transform_cell, cell):
    new_cell = transform_cell(cell)

    assert f"RE({cell.lstrip('< ').strip()})\n" == new_cell
