Further unit tests on SimulatedBroker to increase code coverage.

qstrader/broker/simulated_broker.py

@@ -229,20 +229,21 @@ def create_portfolio(self, portfolio_id, name=None):
         Create a new sub-portfolio with ID 'portfolio_id' and
         an optional name given by 'name'.
         """
-        if portfolio_id in self.portfolios.keys():
+        portfolio_id_str = str(portfolio_id)
+        if portfolio_id_str in self.portfolios.keys():
             raise BrokerException(
                 "Portfolio with ID '%s' already exists. Cannot create "
-                "second portfolio with the same ID." % self.portfolio_id
+                "second portfolio with the same ID." % portfolio_id_str
             )
         else:
             p = Portfolio(
                 self.cur_dt,
                 currency=self.base_currency,
-                portfolio_id=portfolio_id,
+                portfolio_id=portfolio_id_str,
                 name=name
             )
-            self.portfolios[portfolio_id] = p
-            self.open_orders[portfolio_id] = collections.deque()
+            self.portfolios[portfolio_id_str] = p
+            self.open_orders[portfolio_id_str] = collections.deque()
 
     def list_all_portfolios(self):
         """

tests/broker/test_simulated_broker.py

@@ -20,13 +20,15 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import collections
 import unittest
 
 import numpy as np
 import pandas as pd
 import pytz
 
 from qstrader.broker.broker import BrokerException
+from qstrader.broker.portfolio import Portfolio
 from qstrader.broker.simulated_broker import SimulatedBroker
 from qstrader.broker.zero_broker_commission import ZeroBrokerCommission
 from qstrader import settings
@@ -320,6 +322,60 @@ def test_get_account_cash_balance(self):
             sb.get_account_cash_balance(currency="AUD"), 0.0
         )
 
+    def test_create_portfolio(self):
+        """
+        Tests create_portfolio method for:
+        * If portfolio_id already in the dictionary keys,
+        raise BrokerException
+        * If it isn't, check that they portfolio and open
+        orders dictionary was created correctly.
+        """
+        start_dt = pd.Timestamp('2017-10-05 08:00:00', tz=pytz.UTC)
+        exchange = ExchangeMock()
+        sb = SimulatedBroker(start_dt, exchange)
+
+        # If portfolio_id isn't in the dictionary, then check it
+        # was created correctly, along with the orders dictionary
+        sb.create_portfolio(portfolio_id=1234, name="My Portfolio")
+        self.assertTrue("1234" in sb.portfolios)
+        self.assertTrue(isinstance(sb.portfolios["1234"], Portfolio))
+        self.assertTrue("1234" in sb.open_orders)
+        self.assertTrue(isinstance(sb.open_orders["1234"], collections.deque))
+
+        # If portfolio is already in the dictionary
+        # then raise BrokerException
+        with self.assertRaises(BrokerException):
+            sb.create_portfolio(
+                portfolio_id=1234, name="My Portfolio"
+            )
+
+    def test_list_all_portfolio(self):
+        """
+        Tests list_all_portfolios method for:
+        * If empty portfolio dictionary, return empty list
+        * If non-empty, return sorted list via the portfolio IDs
+        """
+        start_dt = pd.Timestamp('2017-10-05 08:00:00', tz=pytz.UTC)
+        exchange = ExchangeMock()
+        sb = SimulatedBroker(start_dt, exchange)
+
+        # If empty portfolio dictionary, return empty list
+        self.assertEquals(sb.list_all_portfolios(), [])
+
+        # If non-empty, return sorted list via the portfolio IDs
+        sb.create_portfolio(portfolio_id=1234, name="My Portfolio #1")
+        sb.create_portfolio(portfolio_id="z154", name="My Portfolio #2")
+        sb.create_portfolio(portfolio_id="abcd", name="My Portfolio #3")
+        self.assertEqual(
+            sorted(
+                [
+                    p.portfolio_id
+                    for p in sb.list_all_portfolios()
+                ]
+            ),
+            ["1234", "abcd", "z154"]
+        )
+
     def test_update_sets_correct_time(self):
         """
         Tests that the update method sets the current