file moves and style changes in sorted/merged layers

tests/test_mongods.py

@@ -0,0 +1,114 @@
+import os
+
+import uuid
+import msgpack
+import pytz
+
+from unittest2 import TestCase
+from pymongo import Connection, ASCENDING
+from itertools import izip, izip_longest
+from datetime import datetime, timedelta
+
+from zipline.gens.mongods import create_pymongo_iterator, MongoTradeHistoryGen
+from zipline.gens.utils import hash_args, assert_datasource_protocol,\
+    assert_trade_protocol, mock_raw_event
+
+import zipline.protocol as zp
+
+mongo_conn_args = {
+    'mongodb_host':  'localhost',
+    'mongodb_port':  27017,
+}
+
+class TempMongo(object):
+
+    def __enter__(self):
+        self.conn = Connection(mongo_conn_args['mongodb_host'],
+                                mongo_conn_args['mongodb_port'])
+
+        temp_id = 'qexec_test_id'
+
+        self.db = self.conn[temp_id]
+
+        return self
+
+    def __exit__(self, type, value, traceback):
+        self.conn.drop_database(self.db.name)
+
+class TestMongoDataGenerator(TestCase):
+    
+    def setUp(self):
+        pass
+    def tearDown(self):
+        pass
+
+    def test_create_pymongo_iterator(self):
+        
+        with TempMongo() as temp_mongo:
+            db = temp_mongo.db
+            coll = db.test
+            coll.ensure_index([('dt', ASCENDING), ('sid', ASCENDING)])
+            
+            for i in xrange(100):
+                # sid = 1, dt ranging from 0 to 99
+                coll.insert(mock_raw_event(1, i))
+            
+            start_date = 20
+            end_date = 50
+            filter = {'sid' : [1]}
+            args = (coll, filter, start_date, end_date)
+            
+            cursor = create_pymongo_iterator(*args)
+            # We filter to only get dt's between 20 and 50
+            expected = (mock_raw_event(1, i) for i in xrange(20, 51))
+            
+            # Assert that our iterator returns the expected values.
+            for cursor_event, expected_event in izip_longest(cursor, expected):
+                del cursor_event['_id']
+                # Easiest way to convert unicode to strings.
+                cursor_event = msgpack.loads(msgpack.dumps(cursor_event))
+                assert expected_event.keys() == cursor_event.keys()
+                assert expected_event.values() == cursor_event.values()
+    
+    def test_MongoTradeHistoryGen(self):
+        
+        with TempMongo() as temp_mongo:
+            db = temp_mongo.db
+            coll = db.test
+            coll.ensure_index([('dt', ASCENDING), ('sid', ASCENDING)])
+            
+            start_date = datetime(year = 2012,month=6,day=5,hour=0)
+            delta = timedelta(hours = 1)
+
+            for i in xrange(100):
+                # sid = 1, dt's increasing an hour at a time from start
+                time = start_date + i * delta
+                coll.insert(mock_raw_event(1, time))
+
+            # Halfway through the events we added to db.
+            end_date = start_date + delta * 50
+
+            filter = {'sid' : [1]}
+            args = (coll, filter, start_date, end_date)
+            db_gen = MongoTradeHistoryGen(*args)
+            
+            expected_times = (start_date + i * delta for i in xrange(51))
+            expected_events = (mock_raw_event(1, t) for t in expected_times)
+            
+            # DB events should match the expected events for price, dt, volume,
+            # and sid. They should also conform to the trade frame protocol.
+
+            for db, expected in izip_longest(db_gen, expected_events):
+                expected['dt'] = expected['dt'].replace(tzinfo = pytz.utc)
+                # Check that our output meets the trade protocol.
+                assert_trade_protocol(db)
+
+                # Check that our output matches expectations
+                for field in iter(['sid', 'dt', 'price', 'volume']):
+                    assert db[field] == expected[field]
+                
+                # Expected output of hash_args:
+                assert db['source_id'] == \
+                    'MongoTradeHistoryGen983a27fd0710414239a5cde71ef5a8fc'
+                
+