"""Tests for core.groups.requires decorator

"""

import numpy as np

import pytest
from MDAnalysis.core.groups import requires
from MDAnalysis import NoDataError

from MDAnalysisTests import make_Universe


class TestRequires(object):
    def test_requires_failure_singular(self):
        @requires("masses")
        def mass_multiplier(ag1, ag2, scalar):
            return (ag1.masses + ag2.masses) * scalar

        u = make_Universe(("charges",))
        with pytest.raises(NoDataError):
            mass_multiplier(u.atoms[:10], u.atoms[20:30], 4.0)

    def test_requires_failure_multiple(self):
        @requires("masses", "charges")
        def mass_multiplier(ag1, ag2, scalar):
            return (ag1.masses + ag2.charges) * scalar

        u = make_Universe(("masses", "types"))
        with pytest.raises(NoDataError):
            mass_multiplier(u.atoms[:10], u.atoms[20:30], 4.0)

    def test_requires_success(self):
        @requires("masses")
        def mass_multiplier(ag1, ag2, scalar):
            return (ag1.masses + ag2.masses) * scalar

        u = make_Universe(("masses",))

        result = mass_multiplier(u.atoms[:10], u.atoms[20:30], 4.0)

        assert isinstance(result, np.ndarray)

    def test_failure_errormessage(self):
        # failures should list all required attributes not
        # just the first one
        @requires("cats", "dogs", "frogs")
        def animal_print(ag):
            return len(ag.cats), len(ag.dogs), len(ag.frogs)

        u = make_Universe()
        try:
            animal_print(u.atoms)
        except NoDataError as e:
            message = e.args[0]
            # Test function name gets returned (for debug)
            assert "animal_print" in message
            assert "cats" in message
            assert "dogs" in message
            assert "frogs" in message
        else:
            pytest.fail(msg="Should raise NoDataError")
