<commit_before>import asyncio
import datetime
from threading import Event

import aioxmpp
import pytest
from pytest import fixture
from asynctest import CoroutineMock, MagicMock

from spade.behaviour import OneShotBehaviour, Behaviour, PeriodicBehaviour, TimeoutBehaviour, FSMBehaviour, State, \
    NotValidState, NotValidTransition, BehaviourNotFinishedException
from spade.message import Message
from spade.template import Template
from tests.utils import make_connected_agent

STATE_ONE = "STATE_ONE"
STATE_TWO = "STATE_TWO"
STATE_THREE = "STATE_THREE"


def wait_for_event(event, tries=100, sleep=0.1):
    counter = 0
    while not event.is_set() and counter < tries:
        event.wait(sleep)
        counter += 1
    if not event.is_set():
        raise Exception("Event was not set")


@fixture
def message():
    return Message(
        to="<EMAIL>",
        sender="<EMAIL>",
        body="message body",
        thread="thread-id",
        metadata={
            "metadata1": "value1",
            "metadata2": "value2"
        }
    )


@fixture
def fsm():
    class StateOne(State):
        async def run(self):
            self.agent.state = STATE_ONE
            self.set_next_state(STATE_TWO)
            self.agent.wait1_behaviour.set()
            await asyncio.sleep(0.01)

    class StateTwo(State):
        async def run(self):
            self.agent.state = STATE_TWO
            self.set_next_state(STATE_THREE)
            self.agent.wait2_behaviour.set()
            await asyncio.sleep(0.01)

    class StateThree(State):
        async def run(self):
            self.agent.state = STATE_THREE
            self.agent.wait3_behaviour.set()

    fsm_ = FSMBehaviour()
    state_one = StateOne()
    state_two = StateTwo()
    state_three = StateThree()
    fsm_.add_state(STATE_ONE, state_one, initial=True)
    fsm_.add_state(STATE_TWO, state_two)
    fsm_.add_state(STATE_THREE, state_three)
    fsm_.add_transition(STATE_ONE, STATE_TWO)
    fsm_.add_transition(STATE_TWO, STATE_THREE)

    return fsm_


def test_on_start_on_end():
    class TestOneShotBehaviour(OneShotBehaviour):
        async def on_start(self):
            self.agent.on_start_flag = True

        async def run(self):
            pass

        async def on_end(self):
            self.agent.on_end_flag = True
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.on_start_flag = False
    agent.on_end_flag = False
    agent.wait_behaviour = Event()
    agent.add_behaviour(TestOneShotBehaviour())

    assert agent.on_start_flag is False
    assert agent.on_end_flag is False

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.on_start_flag is True
    assert agent.on_end_flag is True
    agent.stop()


def test_add_behaviour():
    class EmptyOneShotBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behaviour = EmptyOneShotBehaviour()
    agent.add_behaviour(behaviour)

    assert agent.has_behaviour(behaviour)

    agent.start()

    assert behaviour.agent == agent
    assert behaviour.template is None
    assert behaviour.presence == agent.presence
    assert behaviour.web == agent.web
    assert behaviour.queue.empty()

    wait_for_event(agent.wait_behaviour)

    assert behaviour.done()
    agent.stop()


def test_remove_behaviour():
    class EmptyBehaviour(Behaviour):
        async def run(self):
            pass

    agent = make_connected_agent()
    behaviour = EmptyBehaviour()
    agent.add_behaviour(behaviour)
    assert agent.has_behaviour(behaviour)

    agent.remove_behaviour(behaviour)

    assert not agent.has_behaviour(behaviour)


def test_remove_behaviour_not_added():
    class EmptyBehaviour(Behaviour):
        async def run(self):
            pass

    agent = make_connected_agent()
    behaviour = EmptyBehaviour()

    with pytest.raises(ValueError) as e:
        agent.remove_behaviour(behaviour)


def test_wait_for_agent_start():
    class EmptyOneShotBehaviour(OneShotBehaviour):
        async def on_start(self):
            self.agent.wait_behaviour.set()

        async def run(self):
            pass

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behaviour = EmptyOneShotBehaviour()
    agent.add_behaviour(behaviour)

    assert not agent.wait_behaviour.is_set()
    agent.start()
    wait_for_event(agent.wait_behaviour)
    agent.stop()


def test_behaviour_match():
    class TestBehaviour(OneShotBehaviour):
        async def run(self):
            pass

    template = Template()
    template.sender = "<EMAIL>"
    template.to = "recv1@host"
    template.body = "Hello World"
    template.thread = "thread-id"
    template.metadata = {"performative": "query"}

    behaviour = TestBehaviour()
    behaviour.set_template(template)

    msg = Message()
    msg.sender = "<EMAIL>"
    msg.to = "recv1@host"
    msg.body = "Hello World"
    msg.thread = "thread-id"
    msg.set_metadata("performative", "query")

    assert behaviour.match(msg)


def test_behaviour_match_without_template():
    class TestBehaviour(OneShotBehaviour):
        async def run(self):
            pass

    behaviour = TestBehaviour()

    msg = Message()
    msg.sender = "sender<EMAIL>"
    msg.to = "recv1@host"
    msg.body = "Hello World"
    msg.thread = "thread-id"
    msg.set_metadata("performative", "query")

    assert behaviour.match(msg)


def test_send_message(message):
    class SendBehaviour(OneShotBehaviour):
        async def run(self):
            await self.send(message)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    agent.aiothread.stream = MagicMock()
    agent.stream.send = CoroutineMock()
    agent.add_behaviour(SendBehaviour())
    wait_for_event(agent.wait_behaviour)

    assert agent.stream.send.await_count == 1
    msg_arg = agent.stream.send.await_args[0][0]
    assert msg_arg.body[None] == "message body"
    assert msg_arg.to == aioxmpp.JID.fromstr("to@localhost")
    assert msg_arg.thread == "thread-id"

    agent.stop()


def test_send_message_without_sender():
    class SendBehaviour(OneShotBehaviour):
        async def run(self):
            msg = Message()
            await self.send(msg)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    agent.aiothread.stream = MagicMock()
    agent.stream.send = CoroutineMock()
    agent.add_behaviour(SendBehaviour())
    wait_for_event(agent.wait_behaviour)

    msg_arg = agent.stream.send.await_args[0][0]
    assert msg_arg.from_ == aioxmpp.JID.fromstr("fake@jid")

    agent.stop()


def test_receive():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            msg = Message(body="received body")
            await asyncio.wait_for(self.queue.put(msg), 0.01)
            self.agent.recv_msg = await self.receive()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour)
    assert behaviour.mailbox_size() == 0

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.recv_msg.body == "received body"

    agent.stop()


def test_receive_with_timeout():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.recv_msg = await self.receive(1)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    msg = Message(body="received body")
    template = Template(body="received body")
    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour, template)
    assert behaviour.mailbox_size() == 0

    agent.start()
    agent._message_received(msg.prepare())

    wait_for_event(agent.wait_behaviour)

    assert agent.recv_msg.body == "received body"
    assert agent.recv_msg == msg

    agent.stop()


def test_receive_with_timeout_error():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.recv_msg = await self.receive(0.01)
            agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    template = Template(body="received body")
    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour, template)

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert behaviour.mailbox_size() == 0
    assert agent.recv_msg is None
    agent.stop()


def test_receive_with_empty_queue():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.recv_msg = await self.receive()
            agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    template = Template(body="received body")
    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour, template)

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert behaviour.mailbox_size() == 0
    assert agent.recv_msg is None
    agent.stop()


def test_set_get():
    class SendBehaviour(OneShotBehaviour):
        async def run(self):
            self.set("key", "value")
            assert self.get("key") == "value"
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    agent.add_behaviour(SendBehaviour())
    wait_for_event(agent.wait_behaviour)

    assert agent.get("key") == "value"

    agent.stop()


def test_multiple_templates():
    class Template1Behaviour(OneShotBehaviour):
        async def run(self):
            self.agent.msg1 = await self.receive(timeout=2)

    class Template2Behaviour(OneShotBehaviour):
        async def run(self):
            self.agent.msg2 = await self.receive(timeout=2)

    class Template3Behaviour(OneShotBehaviour):
        async def run(self):
            self.agent.msg3 = await self.receive(timeout=2)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    template1 = Template()
    template1.set_metadata("performative", "template1")
    agent.add_behaviour(Template1Behaviour(), template1)

    template2 = Template()
    template2.set_metadata("performative", "template2")
    agent.add_behaviour(Template2Behaviour(), template2)

    template3 = Template()
    template3.set_metadata("performative", "template3")
    agent.add_behaviour(Template3Behaviour(), template3)

    msg1 = Message(metadata={"performative": "template1"}).prepare()
    msg2 = Message(metadata={"performative": "template2"}).prepare()
    msg3 = Message(metadata={"performative": "template3"}).prepare()

    agent.start()
    agent._message_received(msg1)
    agent._message_received(msg2)
    agent._message_received(msg3)

    wait_for_event(agent.wait_behaviour)

    assert agent.msg1.get_metadata("performative") == "template1"
    assert agent.msg2.get_metadata("performative") == "template2"
    assert agent.msg3.get_metadata("performative") == "template3"
    agent.stop()


def test_kill_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.kill()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()

    agent.add_behaviour(behav)
    wait_for_event(agent.wait_behaviour)

    assert behav.is_killed()
    assert behav.exit_code == 0

    agent.stop()


def test_exit_code_from_kill_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.kill(ValueError)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()

    agent.add_behaviour(behav)
    wait_for_event(agent.wait_behaviour)

    assert behav.is_killed()
    assert behav.exit_code == ValueError

    agent.stop()


def test_set_exit_code_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.exit_code = 1024
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()

    agent.add_behaviour(behav)

    with pytest.raises(BehaviourNotFinishedException):
        assert behav.exit_code

    wait_for_event(agent.wait_behaviour)

    assert not behav.is_killed()
    agent.stop()

    assert behav.exit_code == 1024


def test_notfinishedexception_behaviour():
    class TestBehaviour(OneShotBehaviour):
        async def run(self):
            await self.agent.wait_behaviour.wait()

    agent = make_connected_agent()
    agent.wait_behaviour = asyncio.Event()
    behav = TestBehaviour()
    agent.start()

    agent.add_behaviour(behav)

    with pytest.raises(BehaviourNotFinishedException):
        assert behav.exit_code

    agent.wait_behaviour.set()

    assert behav.exit_code == 0

    agent.stop()


def test_cyclic_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.agent.cycles += 1
            if self.agent.cycles > 2:
                self.kill()
                self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.cycles = 0
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()
    assert agent.cycles == 0

    agent.add_behaviour(behav)
    wait_for_event(agent.wait_behaviour)

    assert agent.cycles == 3
    assert behav.is_killed()

    agent.stop()


def test_oneshot_behaviour():
    class TestOneShotBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.one_shot_behaviour_executed = True
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.one_shot_behaviour_executed = False
    agent.wait_behaviour = Event()
    agent.add_behaviour(TestOneShotBehaviour())

    assert agent.one_shot_behaviour_executed is False

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.one_shot_behaviour_executed is True
    agent.stop()


def test_periodic_behaviour():
    class TestPeriodicBehaviour(PeriodicBehaviour):
        async def run(self):
            self.agent.periodic_behaviour_execution_counter += 1
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.periodic_behaviour_execution_counter = 0
    behaviour = TestPeriodicBehaviour(0.01)
    agent.add_behaviour(behaviour)

    assert agent.periodic_behaviour_execution_counter == 0

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.periodic_behaviour_execution_counter == 1
    agent.stop()


def test_set_period():
    class TestPeriodicBehaviour(PeriodicBehaviour):
        async def run(self):
            pass

    behaviour = TestPeriodicBehaviour(1)
    assert behaviour.period == datetime.timedelta(seconds=1)

    behaviour.period = 2
    assert behaviour.period == datetime.timedelta(seconds=2)


def test_periodic_start_at():
    class TestPeriodicBehaviour(PeriodicBehaviour):
        async def run(self):
            self.agent.delay = datetime.datetime.now()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    start_at = datetime.datetime.now() + datetime.timedelta(seconds=0.01)
    behaviour = TestPeriodicBehaviour(period=0.01, start_at=start_at)

    assert behaviour._next_activation == start_at

    agent.add_behaviour(behaviour)

    wait_for_event(agent.wait_behaviour)
    assert agent.delay >= start_at

    agent.stop()


def test_timeout_behaviour():
    class TestTimeoutBehaviour(TimeoutBehaviour):
        async def run(self):
            self.agent.delay = datetime.datetime.now()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    start_at = datetime.datetime.now() + datetime.timedelta(seconds=0.01)
    behaviour = TestTimeoutBehaviour(start_at=start_at)

    assert behaviour._timeout == start_at

    assert not behaviour._timeout_triggered

    agent.add_behaviour(behaviour)
    wait_for_event(agent.wait_behaviour)

    assert agent.delay >= start_at
    assert behaviour._timeout_triggered
    assert behaviour.done()

    agent.stop()


def test_timeout_behaviour_zero():
    class TestTimeoutBehaviour(TimeoutBehaviour):
        async def run(self):
            self.agent.delay = datetime.datetime.now()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    start_at = datetime.datetime.now() + datetime.timedelta(seconds=0)
    behaviour = TestTimeoutBehaviour(start_at=start_at)

    assert behaviour._timeout == start_at

    assert not behaviour._timeout_triggered

    agent.add_behaviour(behaviour)
    wait_for_event(agent.wait_behaviour)

    assert agent.delay >= start_at
    assert behaviour._timeout_triggered
    assert behaviour.done()

    agent.stop()


def test_fsm_behaviour(fsm):
    agent = make_connected_agent()
    agent.wait1_behaviour = Event()
    agent.wait2_behaviour = Event()
    agent.wait3_behaviour = Event()
    agent.start()

    assert len(fsm._transitions) == 2
    assert fsm.current_state == STATE_ONE

    agent.add_behaviour(fsm)
    wait_for_event(agent.wait1_behaviour)
    assert fsm.current_state == STATE_ONE
    assert agent.state == STATE_ONE

    wait_for_event(agent.wait2_behaviour)
    assert fsm.current_state == STATE_TWO
    assert agent.state == STATE_TWO

    wait_for_event(agent.wait3_behaviour)
    assert fsm.current_state == STATE_THREE
    assert agent.state == STATE_THREE

    agent.stop()


def test_fsm_add_bad_state(fsm):
    with pytest.raises(AttributeError):
        fsm.add_state("BAD_STATE", object())


def test_fsm_is_valid_transition(fsm):
    fsm.is_valid_transition(STATE_ONE, STATE_TWO)


def test_fsm_not_valid_transition(fsm):
    with pytest.raises(NotValidTransition):
        fsm.is_valid_transition(STATE_ONE, STATE_THREE)


def test_fsm_not_valid_state(fsm):
    with pytest.raises(NotValidState):
        fsm.is_valid_transition(STATE_ONE, "BAD_STATE")

    with pytest.raises(NotValidState):
        fsm.is_valid_transition("BAD_STATE", STATE_TWO)


def test_fsm_bad_state():
    class StateOne(State):
        async def run(self):
            self.set_next_state("BAD_STATE")
            self.agent.wait1_behaviour.set()

    class StateTwo(State):
        async def run(self):
            pass

    class BadFSMBehaviour(FSMBehaviour):
        async def on_end(self):
            self.agent.wait2_behaviour.set()

    fsm_ = BadFSMBehaviour()
    state_one = StateOne()
    state_two = StateTwo()
    fsm_.add_state(STATE_ONE, state_one, initial=True)
    fsm_.add_state(STATE_TWO, state_two)
    fsm_.add_transition(STATE_ONE, STATE_TWO)

    agent = make_connected_agent()
    agent.wait1_behaviour = Event()
    agent.wait2_behaviour = Event()
    agent.start()

    assert fsm_.current_state == STATE_ONE

    agent.add_behaviour(fsm_)

    wait_for_event(agent.wait1_behaviour)
    assert fsm_.current_state == STATE_ONE

    wait_for_event(agent.wait2_behaviour)
    assert fsm_.is_killed()

    agent.stop()


def test_fsm_bad_transition():
    class StateOne(State):
        async def run(self):
            self.set_next_state(STATE_THREE)
            self.agent.wait1_behaviour.set()

    class StateTwo(State):
        async def run(self):
            pass

    class StateThree(State):
        async def run(self):
            pass

    class BadFSMBehaviour(FSMBehaviour):
        async def on_end(self):
            self.agent.wait2_behaviour.set()

    fsm_ = BadFSMBehaviour()
    state_one = StateOne()
    state_two = StateTwo()
    state_three = StateThree()
    fsm_.add_state(STATE_ONE, state_one, initial=True)
    fsm_.add_state(STATE_TWO, state_two)
    fsm_.add_state(STATE_THREE, state_three)
    fsm_.add_transition(STATE_ONE, STATE_TWO)
    fsm_.add_transition(STATE_TWO, STATE_THREE)

    agent = make_connected_agent()
    agent.wait1_behaviour = Event()
    agent.wait2_behaviour = Event()
    agent.start()

    assert fsm_.current_state == STATE_ONE

    agent.add_behaviour(fsm_)
    wait_for_event(agent.wait1_behaviour)
    assert fsm_.current_state == STATE_ONE

    wait_for_event(agent.wait2_behaviour)
    assert fsm_.is_killed()

    agent.stop()
<commit_msg>Default tries in test_behaviour set to 1000.<commit_after>import asyncio
import datetime
from threading import Event

import aioxmpp
import pytest
from pytest import fixture
from asynctest import CoroutineMock, MagicMock

from spade.behaviour import OneShotBehaviour, Behaviour, PeriodicBehaviour, TimeoutBehaviour, FSMBehaviour, State, \
    NotValidState, NotValidTransition, BehaviourNotFinishedException
from spade.message import Message
from spade.template import Template
from tests.utils import make_connected_agent

STATE_ONE = "STATE_ONE"
STATE_TWO = "STATE_TWO"
STATE_THREE = "STATE_THREE"


def wait_for_event(event, tries=1000, sleep=0.1):
    counter = 0
    while not event.is_set() and counter < tries:
        event.wait(sleep)
        counter += 1
    if not event.is_set():
        raise Exception("Event was not set")


@fixture
def message():
    return Message(
        to="to@localhost",
        sender="sender@localhost",
        body="message body",
        thread="thread-id",
        metadata={
            "metadata1": "value1",
            "metadata2": "value2"
        }
    )


@fixture
def fsm():
    class StateOne(State):
        async def run(self):
            self.agent.state = STATE_ONE
            self.set_next_state(STATE_TWO)
            self.agent.wait1_behaviour.set()
            await asyncio.sleep(0.01)

    class StateTwo(State):
        async def run(self):
            self.agent.state = STATE_TWO
            self.set_next_state(STATE_THREE)
            self.agent.wait2_behaviour.set()
            await asyncio.sleep(0.01)

    class StateThree(State):
        async def run(self):
            self.agent.state = STATE_THREE
            self.agent.wait3_behaviour.set()

    fsm_ = FSMBehaviour()
    state_one = StateOne()
    state_two = StateTwo()
    state_three = StateThree()
    fsm_.add_state(STATE_ONE, state_one, initial=True)
    fsm_.add_state(STATE_TWO, state_two)
    fsm_.add_state(STATE_THREE, state_three)
    fsm_.add_transition(STATE_ONE, STATE_TWO)
    fsm_.add_transition(STATE_TWO, STATE_THREE)

    return fsm_


def test_on_start_on_end():
    class TestOneShotBehaviour(OneShotBehaviour):
        async def on_start(self):
            self.agent.on_start_flag = True

        async def run(self):
            pass

        async def on_end(self):
            self.agent.on_end_flag = True
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.on_start_flag = False
    agent.on_end_flag = False
    agent.wait_behaviour = Event()
    agent.add_behaviour(TestOneShotBehaviour())

    assert agent.on_start_flag is False
    assert agent.on_end_flag is False

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.on_start_flag is True
    assert agent.on_end_flag is True
    agent.stop()


def test_add_behaviour():
    class EmptyOneShotBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behaviour = EmptyOneShotBehaviour()
    agent.add_behaviour(behaviour)

    assert agent.has_behaviour(behaviour)

    agent.start()

    assert behaviour.agent == agent
    assert behaviour.template is None
    assert behaviour.presence == agent.presence
    assert behaviour.web == agent.web
    assert behaviour.queue.empty()

    wait_for_event(agent.wait_behaviour)

    assert behaviour.done()
    agent.stop()


def test_remove_behaviour():
    class EmptyBehaviour(Behaviour):
        async def run(self):
            pass

    agent = make_connected_agent()
    behaviour = EmptyBehaviour()
    agent.add_behaviour(behaviour)
    assert agent.has_behaviour(behaviour)

    agent.remove_behaviour(behaviour)

    assert not agent.has_behaviour(behaviour)


def test_remove_behaviour_not_added():
    class EmptyBehaviour(Behaviour):
        async def run(self):
            pass

    agent = make_connected_agent()
    behaviour = EmptyBehaviour()

    with pytest.raises(ValueError) as e:
        agent.remove_behaviour(behaviour)


def test_wait_for_agent_start():
    class EmptyOneShotBehaviour(OneShotBehaviour):
        async def on_start(self):
            self.agent.wait_behaviour.set()

        async def run(self):
            pass

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behaviour = EmptyOneShotBehaviour()
    agent.add_behaviour(behaviour)

    assert not agent.wait_behaviour.is_set()
    agent.start()
    wait_for_event(agent.wait_behaviour)
    agent.stop()


def test_behaviour_match():
    class TestBehaviour(OneShotBehaviour):
        async def run(self):
            pass

    template = Template()
    template.sender = "sender1@host"
    template.to = "recv1@host"
    template.body = "Hello World"
    template.thread = "thread-id"
    template.metadata = {"performative": "query"}

    behaviour = TestBehaviour()
    behaviour.set_template(template)

    msg = Message()
    msg.sender = "sender1@host"
    msg.to = "recv1@host"
    msg.body = "Hello World"
    msg.thread = "thread-id"
    msg.set_metadata("performative", "query")

    assert behaviour.match(msg)


def test_behaviour_match_without_template():
    class TestBehaviour(OneShotBehaviour):
        async def run(self):
            pass

    behaviour = TestBehaviour()

    msg = Message()
    msg.sender = "<EMAIL>"
    msg.to = "<EMAIL>"
    msg.body = "Hello World"
    msg.thread = "thread-id"
    msg.set_metadata("performative", "query")

    assert behaviour.match(msg)


def test_send_message(message):
    class SendBehaviour(OneShotBehaviour):
        async def run(self):
            await self.send(message)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    agent.aiothread.stream = MagicMock()
    agent.stream.send = CoroutineMock()
    agent.add_behaviour(SendBehaviour())
    wait_for_event(agent.wait_behaviour)

    assert agent.stream.send.await_count == 1
    msg_arg = agent.stream.send.await_args[0][0]
    assert msg_arg.body[None] == "message body"
    assert msg_arg.to == aioxmpp.JID.fromstr("to@localhost")
    assert msg_arg.thread == "thread-id"

    agent.stop()


def test_send_message_without_sender():
    class SendBehaviour(OneShotBehaviour):
        async def run(self):
            msg = Message()
            await self.send(msg)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    agent.aiothread.stream = MagicMock()
    agent.stream.send = CoroutineMock()
    agent.add_behaviour(SendBehaviour())
    wait_for_event(agent.wait_behaviour)

    msg_arg = agent.stream.send.await_args[0][0]
    assert msg_arg.from_ == aioxmpp.JID.fromstr("fake@jid")

    agent.stop()


def test_receive():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            msg = Message(body="received body")
            await asyncio.wait_for(self.queue.put(msg), 0.01)
            self.agent.recv_msg = await self.receive()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour)
    assert behaviour.mailbox_size() == 0

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.recv_msg.body == "received body"

    agent.stop()


def test_receive_with_timeout():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.recv_msg = await self.receive(1)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    msg = Message(body="received body")
    template = Template(body="received body")
    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour, template)
    assert behaviour.mailbox_size() == 0

    agent.start()
    agent._message_received(msg.prepare())

    wait_for_event(agent.wait_behaviour)

    assert agent.recv_msg.body == "received body"
    assert agent.recv_msg == msg

    agent.stop()


def test_receive_with_timeout_error():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.recv_msg = await self.receive(0.01)
            agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    template = Template(body="received body")
    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour, template)

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert behaviour.mailbox_size() == 0
    assert agent.recv_msg is None
    agent.stop()


def test_receive_with_empty_queue():
    class RecvBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.recv_msg = await self.receive()
            agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    template = Template(body="received body")
    behaviour = RecvBehaviour()
    agent.add_behaviour(behaviour, template)

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert behaviour.mailbox_size() == 0
    assert agent.recv_msg is None
    agent.stop()


def test_set_get():
    class SendBehaviour(OneShotBehaviour):
        async def run(self):
            self.set("key", "value")
            assert self.get("key") == "value"
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    agent.add_behaviour(SendBehaviour())
    wait_for_event(agent.wait_behaviour)

    assert agent.get("key") == "value"

    agent.stop()


def test_multiple_templates():
    class Template1Behaviour(OneShotBehaviour):
        async def run(self):
            self.agent.msg1 = await self.receive(timeout=2)

    class Template2Behaviour(OneShotBehaviour):
        async def run(self):
            self.agent.msg2 = await self.receive(timeout=2)

    class Template3Behaviour(OneShotBehaviour):
        async def run(self):
            self.agent.msg3 = await self.receive(timeout=2)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()

    template1 = Template()
    template1.set_metadata("performative", "template1")
    agent.add_behaviour(Template1Behaviour(), template1)

    template2 = Template()
    template2.set_metadata("performative", "template2")
    agent.add_behaviour(Template2Behaviour(), template2)

    template3 = Template()
    template3.set_metadata("performative", "template3")
    agent.add_behaviour(Template3Behaviour(), template3)

    msg1 = Message(metadata={"performative": "template1"}).prepare()
    msg2 = Message(metadata={"performative": "template2"}).prepare()
    msg3 = Message(metadata={"performative": "template3"}).prepare()

    agent.start()
    agent._message_received(msg1)
    agent._message_received(msg2)
    agent._message_received(msg3)

    wait_for_event(agent.wait_behaviour)

    assert agent.msg1.get_metadata("performative") == "template1"
    assert agent.msg2.get_metadata("performative") == "template2"
    assert agent.msg3.get_metadata("performative") == "template3"
    agent.stop()


def test_kill_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.kill()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()

    agent.add_behaviour(behav)
    wait_for_event(agent.wait_behaviour)

    assert behav.is_killed()
    assert behav.exit_code == 0

    agent.stop()


def test_exit_code_from_kill_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.kill(ValueError)
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()

    agent.add_behaviour(behav)
    wait_for_event(agent.wait_behaviour)

    assert behav.is_killed()
    assert behav.exit_code == ValueError

    agent.stop()


def test_set_exit_code_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.exit_code = 1024
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()

    agent.add_behaviour(behav)

    with pytest.raises(BehaviourNotFinishedException):
        assert behav.exit_code

    wait_for_event(agent.wait_behaviour)

    assert not behav.is_killed()
    agent.stop()

    assert behav.exit_code == 1024


def test_notfinishedexception_behaviour():
    class TestBehaviour(OneShotBehaviour):
        async def run(self):
            await self.agent.wait_behaviour.wait()

    agent = make_connected_agent()
    agent.wait_behaviour = asyncio.Event()
    behav = TestBehaviour()
    agent.start()

    agent.add_behaviour(behav)

    with pytest.raises(BehaviourNotFinishedException):
        assert behav.exit_code

    agent.wait_behaviour.set()

    assert behav.exit_code == 0

    agent.stop()


def test_cyclic_behaviour():
    class TestCyclicBehaviour(Behaviour):
        async def run(self):
            self.agent.cycles += 1
            if self.agent.cycles > 2:
                self.kill()
                self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.cycles = 0
    agent.wait_behaviour = Event()
    behav = TestCyclicBehaviour()
    agent.start()
    assert agent.cycles == 0

    agent.add_behaviour(behav)
    wait_for_event(agent.wait_behaviour)

    assert agent.cycles == 3
    assert behav.is_killed()

    agent.stop()


def test_oneshot_behaviour():
    class TestOneShotBehaviour(OneShotBehaviour):
        async def run(self):
            self.agent.one_shot_behaviour_executed = True
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.one_shot_behaviour_executed = False
    agent.wait_behaviour = Event()
    agent.add_behaviour(TestOneShotBehaviour())

    assert agent.one_shot_behaviour_executed is False

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.one_shot_behaviour_executed is True
    agent.stop()


def test_periodic_behaviour():
    class TestPeriodicBehaviour(PeriodicBehaviour):
        async def run(self):
            self.agent.periodic_behaviour_execution_counter += 1
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.periodic_behaviour_execution_counter = 0
    behaviour = TestPeriodicBehaviour(0.01)
    agent.add_behaviour(behaviour)

    assert agent.periodic_behaviour_execution_counter == 0

    agent.start()
    wait_for_event(agent.wait_behaviour)

    assert agent.periodic_behaviour_execution_counter == 1
    agent.stop()


def test_set_period():
    class TestPeriodicBehaviour(PeriodicBehaviour):
        async def run(self):
            pass

    behaviour = TestPeriodicBehaviour(1)
    assert behaviour.period == datetime.timedelta(seconds=1)

    behaviour.period = 2
    assert behaviour.period == datetime.timedelta(seconds=2)


def test_periodic_start_at():
    class TestPeriodicBehaviour(PeriodicBehaviour):
        async def run(self):
            self.agent.delay = datetime.datetime.now()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    start_at = datetime.datetime.now() + datetime.timedelta(seconds=0.01)
    behaviour = TestPeriodicBehaviour(period=0.01, start_at=start_at)

    assert behaviour._next_activation == start_at

    agent.add_behaviour(behaviour)

    wait_for_event(agent.wait_behaviour)
    assert agent.delay >= start_at

    agent.stop()


def test_timeout_behaviour():
    class TestTimeoutBehaviour(TimeoutBehaviour):
        async def run(self):
            self.agent.delay = datetime.datetime.now()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    start_at = datetime.datetime.now() + datetime.timedelta(seconds=0.01)
    behaviour = TestTimeoutBehaviour(start_at=start_at)

    assert behaviour._timeout == start_at

    assert not behaviour._timeout_triggered

    agent.add_behaviour(behaviour)
    wait_for_event(agent.wait_behaviour)

    assert agent.delay >= start_at
    assert behaviour._timeout_triggered
    assert behaviour.done()

    agent.stop()


def test_timeout_behaviour_zero():
    class TestTimeoutBehaviour(TimeoutBehaviour):
        async def run(self):
            self.agent.delay = datetime.datetime.now()
            self.agent.wait_behaviour.set()

    agent = make_connected_agent()
    agent.wait_behaviour = Event()
    agent.start()

    start_at = datetime.datetime.now() + datetime.timedelta(seconds=0)
    behaviour = TestTimeoutBehaviour(start_at=start_at)

    assert behaviour._timeout == start_at

    assert not behaviour._timeout_triggered

    agent.add_behaviour(behaviour)
    wait_for_event(agent.wait_behaviour)

    assert agent.delay >= start_at
    assert behaviour._timeout_triggered
    assert behaviour.done()

    agent.stop()


def test_fsm_behaviour(fsm):
    agent = make_connected_agent()
    agent.wait1_behaviour = Event()
    agent.wait2_behaviour = Event()
    agent.wait3_behaviour = Event()
    agent.start()

    assert len(fsm._transitions) == 2
    assert fsm.current_state == STATE_ONE

    agent.add_behaviour(fsm)
    wait_for_event(agent.wait1_behaviour)
    assert fsm.current_state == STATE_ONE
    assert agent.state == STATE_ONE

    wait_for_event(agent.wait2_behaviour)
    assert fsm.current_state == STATE_TWO
    assert agent.state == STATE_TWO

    wait_for_event(agent.wait3_behaviour)
    assert fsm.current_state == STATE_THREE
    assert agent.state == STATE_THREE

    agent.stop()


def test_fsm_add_bad_state(fsm):
    with pytest.raises(AttributeError):
        fsm.add_state("BAD_STATE", object())


def test_fsm_is_valid_transition(fsm):
    fsm.is_valid_transition(STATE_ONE, STATE_TWO)


def test_fsm_not_valid_transition(fsm):
    with pytest.raises(NotValidTransition):
        fsm.is_valid_transition(STATE_ONE, STATE_THREE)


def test_fsm_not_valid_state(fsm):
    with pytest.raises(NotValidState):
        fsm.is_valid_transition(STATE_ONE, "BAD_STATE")

    with pytest.raises(NotValidState):
        fsm.is_valid_transition("BAD_STATE", STATE_TWO)


def test_fsm_bad_state():
    class StateOne(State):
        async def run(self):
            self.set_next_state("BAD_STATE")
            self.agent.wait1_behaviour.set()

    class StateTwo(State):
        async def run(self):
            pass

    class BadFSMBehaviour(FSMBehaviour):
        async def on_end(self):
            self.agent.wait2_behaviour.set()

    fsm_ = BadFSMBehaviour()
    state_one = StateOne()
    state_two = StateTwo()
    fsm_.add_state(STATE_ONE, state_one, initial=True)
    fsm_.add_state(STATE_TWO, state_two)
    fsm_.add_transition(STATE_ONE, STATE_TWO)

    agent = make_connected_agent()
    agent.wait1_behaviour = Event()
    agent.wait2_behaviour = Event()
    agent.start()

    assert fsm_.current_state == STATE_ONE

    agent.add_behaviour(fsm_)

    wait_for_event(agent.wait1_behaviour)
    assert fsm_.current_state == STATE_ONE

    wait_for_event(agent.wait2_behaviour)
    assert fsm_.is_killed()

    agent.stop()


def test_fsm_bad_transition():
    class StateOne(State):
        async def run(self):
            self.set_next_state(STATE_THREE)
            self.agent.wait1_behaviour.set()

    class StateTwo(State):
        async def run(self):
            pass

    class StateThree(State):
        async def run(self):
            pass

    class BadFSMBehaviour(FSMBehaviour):
        async def on_end(self):
            self.agent.wait2_behaviour.set()

    fsm_ = BadFSMBehaviour()
    state_one = StateOne()
    state_two = StateTwo()
    state_three = StateThree()
    fsm_.add_state(STATE_ONE, state_one, initial=True)
    fsm_.add_state(STATE_TWO, state_two)
    fsm_.add_state(STATE_THREE, state_three)
    fsm_.add_transition(STATE_ONE, STATE_TWO)
    fsm_.add_transition(STATE_TWO, STATE_THREE)

    agent = make_connected_agent()
    agent.wait1_behaviour = Event()
    agent.wait2_behaviour = Event()
    agent.start()

    assert fsm_.current_state == STATE_ONE

    agent.add_behaviour(fsm_)
    wait_for_event(agent.wait1_behaviour)
    assert fsm_.current_state == STATE_ONE

    wait_for_event(agent.wait2_behaviour)
    assert fsm_.is_killed()

    agent.stop()

{'fraction_non_alphanumeric': 0.09892566154733579, 'fraction_numerical': 0.006537361501670392, 'mean_word_length': 3.181708542713568, 'pattern_counts': {'":': 12, '<': 12, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2724616', 'n_tokens_mistral': 14010, 'n_tokens_neox': 12959, 'n_words': 2308}