lemur/lemur/tests/test_schemas.py

import pytest
from marshmallow.exceptions import ValidationError

from lemur.tests.factories import RoleFactory


def test_get_object_attribute():
    from lemur.schemas import get_object_attribute

    with pytest.raises(ValidationError):
        get_object_attribute({})

    with pytest.raises(ValidationError):
        get_object_attribute([{}], many=True)

    with pytest.raises(ValidationError):
        get_object_attribute([{}, {"id": 1}], many=True)

    with pytest.raises(ValidationError):
        get_object_attribute([{}, {"name": "test"}], many=True)

    assert get_object_attribute({"name": "test"}) == "name"
    assert get_object_attribute({"id": 1}) == "id"
    assert get_object_attribute([{"name": "test"}], many=True) == "name"
    assert get_object_attribute([{"id": 1}], many=True) == "id"


def test_fetch_objects(session):
    from lemur.roles.models import Role
    from lemur.schemas import fetch_objects

    role = RoleFactory()
    role1 = RoleFactory()
    session.commit()

    data = {"id": role.id}
    found_role = fetch_objects(Role, data)
    assert found_role == role

    data = {"name": role.name}
    found_role = fetch_objects(Role, data)
    assert found_role == role

    data = [{"id": role.id}, {"id": role1.id}]
    found_roles = fetch_objects(Role, data, many=True)
    assert found_roles == [role, role1]

    data = [{"name": role.name}, {"name": role1.name}]
    found_roles = fetch_objects(Role, data, many=True)
    assert found_roles == [role, role1]

    with pytest.raises(ValidationError):
        data = [{"name": "blah"}, {"name": role1.name}]
        fetch_objects(Role, data, many=True)

    with pytest.raises(ValidationError):
        data = {"name": "nah"}
        fetch_objects(Role, data)