TST: adds tmp_assets_db and tmp_asset_finder so we don't need an entire trading env

zipline/utils/test_utils.py

@@ -2,21 +2,26 @@ from contextlib import contextmanager
 from itertools import (
     product,
 )
+import operator
+import os
+import shutil
+import tempfile
+
 from logbook import FileHandler
 from mock import patch
 from numpy.testing import assert_allclose, assert_array_equal
-import operator
-from zipline.finance.blotter import ORDER_STATUS
-from zipline.utils import security_list
+import pandas as pd
 from six import (
     itervalues,
 )
 from six.moves import filter
+from sqlalchemy import create_engine
+
+from zipline.assets import AssetFinder
+from zipline.assets.asset_writer import AssetDBWriterFromDictionary
+from zipline.finance.blotter import ORDER_STATUS
+from zipline.utils import security_list
 
-import os
-import pandas as pd
-import shutil
-import tempfile
 
 EPOCH = pd.Timestamp(0, tz='UTC')
 
@@ -356,3 +361,57 @@ class ExplodingObject(object):
     """
     def __getattribute__(self, name):
         raise UnexpectedAttributeAccess(name)
+
+
+class tmp_assets_db(object):
+    """Create a temporary assets sqlite database.
+    This is meant to be used as a context manager.
+
+    Paramaters
+    ----------
+    data : dict, optional
+        The data to feed to the writer. By default this maps:
+        ('A', 'B', 'C') -> map(ord, 'ABC')
+    """
+    def __init__(self, data=None):
+        self._eng = None
+        self._data = AssetDBWriterFromDictionary(
+            data if data else {
+                ord('A'): {
+                    'symbol': 'A',
+                    'start_date': pd.Timestamp(0),
+                    'end_date': pd.Timestamp('2015'),
+                },
+                ord('B'): {
+                    'symbol': 'B',
+                    'start_date': pd.Timestamp(0),
+                    'end_date': pd.Timestamp('2015'),
+                },
+                ord('C'): {
+                    'symbol': 'C',
+                    'start_date': pd.Timestamp(0),
+                    'end_date': pd.Timestamp('2015'),
+                },
+            },
+        )
+
+    def __enter__(self):
+        self._eng = eng = create_engine('sqlite://')
+        self._data.write_all(eng)
+        return eng
+
+    def __exit__(self, *excinfo):
+        assert self._eng is not None, '_eng was not set in __enter__'
+        self._eng.dispose()
+
+
+class tmp_asset_finder(tmp_assets_db):
+    """Create a temporary asset finder using an in memory sqlite db.
+
+    Paramaters
+    ----------
+    data : dict, optional
+        The data to feed to the writer
+    """
+    def __enter__(self):
+        return AssetFinder(super(tmp_asset_finder, self).__enter__())