import numpy as np
import pytest

from mbtrack2 import Optics


def test_local_optics():
    beta = np.array([1, 1])
    alpha = np.array([0, 0])
    dispersion = np.array([0, 0, 0, 0])
    gamma = (1 + alpha**2) / beta
    optics = Optics(local_beta=beta, local_alpha=alpha, 
                      local_dispersion=dispersion)
    np.testing.assert_allclose(optics.local_beta, beta)
    np.testing.assert_allclose(optics.local_alpha, alpha)
    np.testing.assert_allclose(optics.local_dispersion, dispersion)
    np.testing.assert_allclose(optics.local_gamma, gamma)