import numpy as np
import pytest

from latdraw import elements

LENGTH = 0.5


def test_init_Drift():
    name = "drift"
    drift = elements.Drift(name, LENGTH)

    assert drift.name == name
    assert drift.length == LENGTH


def test_init_SBend():
    name = "sbend"
    angle = np.pi / 10
    sbend = elements.SBend(name, LENGTH, angle)

    assert sbend.name == name
    assert sbend.length == LENGTH
    assert sbend.angle == angle


def test_init_RBend():
    name = "rbend"
    angle = np.pi / 10
    sbend = elements.RBend(name, LENGTH, angle)

    assert sbend.name == name
    assert sbend.length == LENGTH
    assert sbend.angle == angle


def test_init_Quadrupole():
    name = "quadrupole"
    k1 = 0.005
    sbend = elements.Quadrupole(name, LENGTH, k1)

    assert sbend.name == name
    assert sbend.length == LENGTH
    assert sbend.k1 == k1


def test_init_Sextupole():
    name = "sextupole"
    k2 = 0.005
    sbend = elements.Sextupole(name, LENGTH, k2)

    assert sbend.name == name
    assert sbend.length == LENGTH
    assert sbend.k2 == k2


def test_init_Octupole():
    name = "octupole"
    k3 = 0.005
    sbend = elements.Octupole(name, LENGTH, k3)

    assert sbend.name == name
    assert sbend.length == LENGTH
    assert sbend.k3 == k3


def test_init_VKicker():
    name = "vkicker"
    angle = 0.005
    sbend = elements.VKicker(name, LENGTH, angle)

    assert sbend.name == name
    assert sbend.length == LENGTH
    assert sbend.angle == angle


def test_init_HKicker():
    name = "hkicker"
    angle = 0.005
    sbend = elements.HKicker(name, LENGTH, angle)

    assert sbend.name == name
    assert sbend.length == LENGTH
    assert sbend.angle == angle
