diff --git a/axelrod/action.py b/axelrod/action.py index 901723ce1..b2bd8c1e9 100644 --- a/axelrod/action.py +++ b/axelrod/action.py @@ -7,6 +7,7 @@ """ from enum import Enum +from functools import total_ordering from typing import Iterable @@ -17,6 +18,7 @@ def __init__(self, *args): super(UnknownActionError, self).__init__(*args) +@total_ordering class Action(Enum): """Core actions in the Prisoner's Dilemna. @@ -30,6 +32,15 @@ class Action(Enum): def __bool__(self): return bool(self.value) + def __eq__(self, other): + return self.value == other.value + + def __hash__(self): + return hash(self.value) + + def __lt__(self, other): + return self.value < other.value + def __repr__(self): return "{}".format(self.name) diff --git a/axelrod/tests/strategies/test_sequence_player.py b/axelrod/tests/strategies/test_sequence_player.py index a313d98ec..429a4e9f4 100644 --- a/axelrod/tests/strategies/test_sequence_player.py +++ b/axelrod/tests/strategies/test_sequence_player.py @@ -22,8 +22,7 @@ def test_sequence_player(self): """Basic test for SequencePlayer.""" def cooperate_gen(): - yield C - + yield 1 player = SequencePlayer(generator_function=cooperate_gen) opponent = TestOpponent() self.assertEqual(C, player.strategy(opponent)) diff --git a/axelrod/tests/unit/test_actions.py b/axelrod/tests/unit/test_actions.py index 4f07f17cd..a7e1a0e9a 100644 --- a/axelrod/tests/unit/test_actions.py +++ b/axelrod/tests/unit/test_actions.py @@ -7,6 +7,17 @@ class TestAction(unittest.TestCase): + def test_eq(self): + self.assertEqual(C, C) + self.assertEqual(D, D) + + def test_hash(self): + self.assertEqual(hash(C), 1) + self.assertEqual(hash(D), 0) + + def test_lt(self): + self.assertLess(D, C) + def test_repr(self): self.assertEqual(repr(C), "C") self.assertEqual(repr(D), "D")