import math

from src.compute import rot_2D, rot_3D

def test_get_2D_rotation_matrix():

    #Arrange
    mat = [[0], [0]]
    angle = math.pi
    expected_matrix = [[math.cos(math.pi), -math.sin(math.pi)], [math.sin(math.pi), math.cos(math.pi)]]

    #Act
    rot = rot_2D(mat, angle)

    #Assert
    assert rot == expected_matrix

def test_get_3D_rotation_matrix():

    #Arrange
    angle = math.pi
    axis = "x"
    expected_matrix = [
        [1, 0, 0],
        [0, math.cos(math.pi), -math.sin(math.pi)],
        [0, math.sin(math.pi), math.cos(math.pi)]]

    #Act
    rot = rot_3D(angle, axis)

    #Assert
    assert rot == expected_matrix