You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
soil/tests/test_agents.py

525 lines
16 KiB
Python

from unittest import TestCase
import pytest
from soil import agents, events, environment
from soil import time as stime
class Dead(agents.FSM):
@agents.default_state
@agents.state
def only(self):
return self.die()
class TestAgents(TestCase):
def test_die_returns_infinity(self):
"""The last step of a dead agent should return time.INFINITY"""
d = Dead(unique_id=0, model=environment.Environment())
ret = d.step()
assert ret == stime.INFINITY
def test_die_raises_exception(self):
"""A dead agent should continue returning INFINITY after death"""
d = Dead(unique_id=0, model=environment.Environment())
assert d.alive
d.step()
assert not d.alive
when = float(d.step())
assert not d.alive
assert when == stime.INFINITY
def test_agent_generator(self):
"""
The step function of an agent could be a generator. In that case, the state of the
agent will be resumed after every call to step.
"""
a = 0
class Gen(agents.BaseAgent):
def step(self):
nonlocal a
for i in range(5):
yield
a += 1
e = environment.Environment()
g = Gen(model=e, unique_id=e.next_id())
e.schedule.add(g)
for i in range(5):
e.step()
assert a == i
def test_state_decorator(self):
class MyAgent(agents.FSM):
times_run = 0
@agents.state("original", default=True)
def root(self):
return self.other
@agents.state
def other(self):
self.times_run += 1
assert MyAgent.other.id == "other"
e = environment.Environment()
a = e.add_agent(MyAgent)
e.step()
assert a.times_run == 0
a.step()
assert a.times_run == 1
assert a.state_id == MyAgent.other.id
a.step()
assert a.times_run == 2
def test_state_decorator_multiple(self):
class MyAgent(agents.FSM):
times_run = 0
@agents.state(default=True)
def one(self):
return self.two
@agents.state
def two(self):
return self.one
e = environment.Environment()
first = e.add_agent(MyAgent, state_id=MyAgent.one)
second = e.add_agent(MyAgent, state_id=MyAgent.two)
assert first.state_id == MyAgent.one.id
assert second.state_id == MyAgent.two.id
e.step()
assert first.state_id == MyAgent.two.id
assert second.state_id == MyAgent.one.id
def test_state_decorator_multiple_async(self):
class MyAgent(agents.FSM):
times_run = 0
@agents.state(default=True)
def one(self):
yield self.delay(1)
return self.two
@agents.state
def two(self):
yield self.delay(1)
return self.one
e = environment.Environment()
first = e.add_agent(MyAgent, state_id=MyAgent.one)
second = e.add_agent(MyAgent, state_id=MyAgent.two)
for i in range(2):
assert first.state_id == MyAgent.one.id
assert second.state_id == MyAgent.two.id
e.step()
for i in range(2):
assert first.state_id == MyAgent.two.id
assert second.state_id == MyAgent.one.id
e.step()
def test_broadcast(self):
"""
An agent should be able to broadcast messages to every other agent, AND each receiver should be able
to process it
"""
class BCast(agents.Evented):
pings_received = []
async def step(self):
self.broadcast("PING")
print("PING sent")
while True:
msgs = await self.received()
self.pings_received += msgs
e = environment.Environment()
num_agents = 10
for i in range(num_agents):
e.add_agent(agent_class=BCast)
e.step()
# Agents are executed in order, so the first agent should have not received any messages
pings_received = lambda: [len(a.pings_received) for a in e.agents]
assert sorted(pings_received()) == list(range(0, num_agents))
e.step()
# After the second step, every agent should have received a broadcast from every other agent
received = pings_received()
assert all(x == (num_agents - 1) for x in received)
def test_ask_messages(self):
"""
An agent should be able to ask another agent, and wait for a response.
"""
# There are two agents, they try to send pings
# This is arguably a very contrived example.
# There should be a delay of one step between agent 0 and 1
# On the first step:
# Agent 0 sends a PING, but blocks before a PONG
# Agent 1 detects the PING, responds with a PONG, and blocks after its own PING
# After that step, every agent can both receive (there are pending messages) and send.
# In each step, for each agent, one message is sent, and another one is received
# (although not necessarily in that order).
# Results depend on ordering (agents are normally shuffled)
# so we force the timedactivation not to be shuffled
pings = []
pongs = []
responses = []
class Ping(agents.EventedAgent):
def step(self):
target_id = (self.unique_id + 1) % self.count_agents()
target = self.model.agents[target_id]
print("starting")
while True:
if pongs or not pings: # First agent, or anyone after that
pings.append(self.now)
response = yield from target.ask("PING")
responses.append(response)
else:
print("NOT sending ping")
print("Checking msgs")
# Do not block if we have already received a PING
msgs = yield from self.received()
for ping in msgs:
if ping.payload == "PING":
ping.reply = "PONG"
pongs.append(self.now)
else:
raise Exception("This should never happen")
e = environment.Environment(schedule_class=stime.OrderedTimedActivation)
for i in range(2):
e.add_agent(agent_class=Ping)
assert e.now == 0
for i in range(5):
e.step()
time = i + 1
assert e.now == time
assert len(pings) == 2 * time
assert len(pongs) == (2 * time) - 1
# Every step between 0 and t appears twice
assert sum(pings) == sum(range(time)) * 2
# It is the same as pings, without the leading 0
assert sum(pongs) == sum(range(time)) * 2
def test_agent_filter(self):
e = environment.Environment()
e.add_agent(agent_class=agents.BaseAgent)
e.add_agent(agent_class=agents.Evented)
base = list(e.agents(agent_class=agents.BaseAgent))
assert len(base) == 2
ev = list(e.agents(agent_class=agents.Evented))
assert len(ev) == 1
assert ev[0].unique_id == 1
null = list(e.agents(unique_ids=[0, 1], agent_class=agents.NetworkAgent))
assert not null
def test_agent_return(self):
'''
An agent should be able to cycle through different states and control when it
should be awaken.
'''
class TestAgent(agents.Agent):
@agents.state(default=True)
def one(self):
return self.two
@agents.state
def two(self):
return self.three.at(10)
@agents.state
def three(self):
return self.four.delay(1)
@agents.state
def four(self):
yield self.delay(2)
return self.five.delay(3)
@agents.state
def five(self):
return self.delay(1)
model = environment.Environment()
a = model.add_agent(TestAgent)
assert a.state_id == TestAgent.one.id
assert a.now == 0
model.step()
assert a.state_id == TestAgent.two.id
assert a.now == 1
model.step()
assert a.state_id == TestAgent.three.id
assert a.now == 10
model.step()
assert a.state_id == TestAgent.four.id
assert a.now == 11
model.step()
assert a.state_id == TestAgent.four.id
assert a.now == 13
model.step()
assert a.state_id == TestAgent.five.id
assert a.now == 16
model.step()
assert a.state_id == TestAgent.five.id
assert a.now == 17
def test_agent_async(self):
'''
Async functions should also be valid states.
'''
class TestAgent(agents.Agent):
@agents.state(default=True)
def one(self):
return self.two
@agents.state
def two(self):
return self.three.at(10)
@agents.state
def three(self):
return self.four.delay(1)
@agents.state
async def four(self):
await self.delay(2)
return self.five.delay(3)
@agents.state
def five(self):
return self.delay(1)
model = environment.Environment()
a = model.add_agent(TestAgent)
assert a.now == 0
assert a.state_id == TestAgent.one.id
model.step()
assert a.now == 1
assert a.state_id == TestAgent.two.id
model.step()
assert a.now == 10
assert a.state_id == TestAgent.three.id
model.step()
assert a.state_id == TestAgent.four.id
assert a.now == 11
model.step()
assert a.state_id == TestAgent.four.id
assert a.now == 13
model.step()
assert a.state_id == TestAgent.five.id
assert a.now == 16
model.step()
assert a.state_id == TestAgent.five.id
assert a.now == 17
def test_agent_return_step(self):
'''
The same result as the previous test should be achievable by manually
handling the agent state.
'''
class TestAgent(agents.Agent):
my_state = 1
my_count = 0
def step(self):
if self.my_state == 1:
self.my_state = 2
return None
elif self.my_state == 2:
self.my_state = 3
return self.at(10)
elif self.my_state == 3:
self.my_state = 4
self.my_count = 0
return self.delay(1)
elif self.my_state == 4:
self.my_count += 1
if self.my_count == 1:
return self.delay(2)
self.my_state = 5
return self.delay(3)
elif self.my_state == 5:
return self.delay(1)
model = environment.Environment()
a = model.add_agent(TestAgent)
assert a.my_state == 1
assert a.now == 0
model.step()
assert a.now == 1
assert a.my_state == 2
model.step()
assert a.now == 10
assert a.my_state == 3
model.step()
assert a.now == 11
assert a.my_state == 4
model.step()
assert a.now == 13
assert a.my_state == 4
model.step()
assert a.now == 16
assert a.my_state == 5
model.step()
assert a.now == 17
assert a.my_state == 5
def test_agent_return_step_async(self):
'''
The same result as the previous test should be achievable by manually
handling the agent state.
'''
class TestAgent(agents.Agent):
my_state = 1
async def step(self):
self.my_state = 2
await self.delay()
self.my_state = 3
await self.at(10)
self.my_state = 4
await self.delay(1)
await self.delay(2)
self.my_state = 5
await self.delay(3)
while True:
await self.delay(1)
model = environment.Environment()
a = model.add_agent(TestAgent)
assert a.my_state == 1
assert a.now == 0
model.step()
assert a.now == 1
assert a.my_state == 2
model.step()
assert a.now == 10
assert a.my_state == 3
model.step()
assert a.now == 11
assert a.my_state == 4
model.step()
assert a.now == 13
assert a.my_state == 4
model.step()
assert a.now == 16
assert a.my_state == 5
model.step()
assert a.now == 17
assert a.my_state == 5
def test_receive(self):
'''
An agent should be able to receive a message after waiting
'''
model = environment.Environment()
class TestAgent(agents.Agent):
sent = False
woken = 0
def step(self):
self.woken += 1
return super().step()
@agents.state(default=True)
async def one(self):
try:
self.sent = await self.received(timeout=15)
return self.two.at(20)
except events.TimedOut:
pass
@agents.state
def two(self):
return self.die()
a = model.add_agent(TestAgent)
class Sender(agents.Agent):
async def step(self):
await self.delay(10)
a.tell(1)
return stime.INFINITY
b = model.add_agent(Sender)
# Start and wait
model.step()
assert model.now == 10
assert a.woken == 1
assert not a.sent
# Sending the message
model.step()
assert model.now == 10
assert a.woken == 1
assert not a.sent
# The receiver callback
model.step()
assert model.now == 15
assert a.woken == 2
assert a.sent[0].payload == 1
# The timeout
model.step()
assert model.now == 20
assert a.woken == 2
# The last state of the agent
model.step()
assert a.woken == 3
assert model.now == float('inf')
def test_receive_timeout(self):
'''
A timeout should be raised if no messages are received after an expiration time
'''
model = environment.Environment()
timedout = False
class TestAgent(agents.Agent):
@agents.state(default=True)
def one(self):
try:
yield from self.received(timeout=10)
raise AssertionError('Should have raised an error.')
except events.TimedOut:
nonlocal timedout
timedout = True
a = model.add_agent(TestAgent)
model.step()
assert model.now == 10
model.step()
# Wake up the callback
assert model.now == 10
assert not timedout
# The actual timeout
model.step()
assert model.now == 11
assert timedout
def test_attributes(self):
"""Attributes should be individual per agent"""
class MyAgent(agents.Agent):
my_attribute = 0
model = environment.Environment()
a = MyAgent(model=model)
assert a.my_attribute == 0
b = MyAgent(model=model, my_attribute=1)
assert b.my_attribute == 1
assert a.my_attribute == 0