Spaces:
Sleeping
Sleeping
File size: 1,507 Bytes
c71bf62 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | import unittest
from dispatch_arena.models import (
Action,
Config,
Courier,
DispatchArenaObservation,
Mode,
Observation,
Order,
OrderStatus,
State,
)
class DispatchArenaModelsTests(unittest.TestCase):
def test_action_roundtrip(self):
action = Action(action_type="assign", courier_id="courier_0", order_id="order_0")
self.assertEqual(action.name, "assign")
self.assertEqual(Action.from_dict(action.to_dict()), action)
def test_config_clamps_normal_size(self):
config = Config(mode=Mode.NORMAL, num_couriers=9, num_orders=99)
self.assertEqual(config.num_couriers, 5)
self.assertEqual(config.num_orders, 10)
def test_state_serializes_public_fields_without_hidden_prep_when_none(self):
state = State(
orders=[
Order(
id="order_0",
pickup_node_id="pickup",
dropoff_node_id="dropoff",
status=OrderStatus.READY,
prep_remaining=None,
)
],
couriers=[Courier(id="courier_0", node_id="hub")],
)
data = state.to_dict()
self.assertEqual(data["orders"][0]["status"], "ready")
self.assertNotIn("prep_remaining", str(data))
def test_observation_alias_exports(self):
obs = Observation()
self.assertIsInstance(obs, DispatchArenaObservation)
if __name__ == "__main__":
unittest.main()
|