Update tests and add docstring

axelrod/moran.py

@@ -387,6 +387,21 @@ def __len__(self) -> int:
         return len(self.populations)
 
     def populations_plot(self, ax=None):
+        """
+        Create a stackplot of the population distributions at each iteration of
+        the Moran process.
+
+        Parameters
+        ----------------
+        ax: matplotlib axis
+            Allows the plot to be written to a given matplotlib axis.
+            Default is None.
+
+        Returns
+        -----------
+        A matplotlib axis object
+
+        """
         player_names = self.populations[0].keys()
         if ax is None:
             _, ax = plt.subplots()

axelrod/tests/unit/test_moran.py

@@ -314,7 +314,6 @@ def test_constant_fitness_case(self):
         winners = Counter(winners)
         self.assertEqual(winners["Defector"], 88)
 
-
     def test_cache(self):
         p1, p2 = axelrod.Cooperator(), axelrod.Defector()
         mp = MoranProcess((p1, p2))
@@ -334,17 +333,18 @@ def test_iter(self):
     def test_population_plot(self):
         # Test that can plot on a given matplotlib axes
         axelrod.seed(15)
-        N = 5
-        players = [random.choice(axelrod.basic_strategies)() for _ in range(N)]
-        mp = axelrod.MoranProcess(players=players, turns=200)
+        players = [random.choice(axelrod.demo_strategies)() for _ in range(5)]
+        mp = axelrod.MoranProcess(players=players, turns=30)
         mp.play()
         fig, axarr = plt.subplots(2, 2)
         ax = axarr[1, 0]
         mp.populations_plot(ax=ax)
-
         self.assertEqual(ax.get_xlim(), (-0.8, 16.8))
         self.assertEqual(ax.get_ylim(), (0, 5.25))
 
+        # Run with a give axis
+        mp.populations_plot()
+
 
 class GraphMoranProcess(unittest.TestCase):
 

docs/tutorials/getting_started/moran.rst

@@ -84,12 +84,12 @@ use a larger population to get a bit more data::
    >>> import matplotlib.pyplot as plt
    >>> axl.seed(15)  # for reproducible example
    >>> N = 10
-   >>> players = [random.choice(axl.basic_strategies)() for i in range(N)]
+   >>> players = [random.choice(axl.demo_strategies)() for i in range(N)]
    >>> mp = axl.MoranProcess(players=players, turns=200)
    >>> populations = mp.play()
    >>> mp.winning_strategy_name
-   'Tit For Tat'
-    >>> ax = mp.populations_plot()
+   'Defector'
+   >>> ax = mp.populations_plot()
    >>> plt.show()  #doctest: +SKIP