Show More
@@ -1,224 +1,233 b'' | |||||
1 | """Tests for db backends |
|
1 | """Tests for db backends | |
2 |
|
2 | |||
3 | Authors: |
|
3 | Authors: | |
4 |
|
4 | |||
5 | * Min RK |
|
5 | * Min RK | |
6 | """ |
|
6 | """ | |
7 |
|
7 | |||
8 | #------------------------------------------------------------------------------- |
|
8 | #------------------------------------------------------------------------------- | |
9 | # Copyright (C) 2011 The IPython Development Team |
|
9 | # Copyright (C) 2011 The IPython Development Team | |
10 | # |
|
10 | # | |
11 | # Distributed under the terms of the BSD License. The full license is in |
|
11 | # Distributed under the terms of the BSD License. The full license is in | |
12 | # the file COPYING, distributed as part of this software. |
|
12 | # the file COPYING, distributed as part of this software. | |
13 | #------------------------------------------------------------------------------- |
|
13 | #------------------------------------------------------------------------------- | |
14 |
|
14 | |||
15 | #------------------------------------------------------------------------------- |
|
15 | #------------------------------------------------------------------------------- | |
16 | # Imports |
|
16 | # Imports | |
17 | #------------------------------------------------------------------------------- |
|
17 | #------------------------------------------------------------------------------- | |
18 |
|
18 | |||
19 | from __future__ import division |
|
19 | from __future__ import division | |
20 |
|
20 | |||
|
21 | import os | |||
21 | import tempfile |
|
22 | import tempfile | |
22 | import time |
|
23 | import time | |
23 |
|
24 | |||
24 | from datetime import datetime, timedelta |
|
25 | from datetime import datetime, timedelta | |
25 | from unittest import TestCase |
|
26 | from unittest import TestCase | |
26 |
|
27 | |||
27 | from IPython.parallel import error |
|
28 | from IPython.parallel import error | |
28 | from IPython.parallel.controller.dictdb import DictDB |
|
29 | from IPython.parallel.controller.dictdb import DictDB | |
29 | from IPython.parallel.controller.sqlitedb import SQLiteDB |
|
30 | from IPython.parallel.controller.sqlitedb import SQLiteDB | |
30 | from IPython.parallel.controller.hub import init_record, empty_record |
|
31 | from IPython.parallel.controller.hub import init_record, empty_record | |
31 |
|
32 | |||
32 | from IPython.testing import decorators as dec |
|
33 | from IPython.testing import decorators as dec | |
33 | from IPython.zmq.session import Session |
|
34 | from IPython.zmq.session import Session | |
34 |
|
35 | |||
35 |
|
36 | |||
36 | #------------------------------------------------------------------------------- |
|
37 | #------------------------------------------------------------------------------- | |
37 | # TestCases |
|
38 | # TestCases | |
38 | #------------------------------------------------------------------------------- |
|
39 | #------------------------------------------------------------------------------- | |
39 |
|
40 | |||
40 | class TestDictBackend(TestCase): |
|
41 | class TestDictBackend(TestCase): | |
41 | def setUp(self): |
|
42 | def setUp(self): | |
42 | self.session = Session() |
|
43 | self.session = Session() | |
43 | self.db = self.create_db() |
|
44 | self.db = self.create_db() | |
44 | self.load_records(16) |
|
45 | self.load_records(16) | |
45 |
|
46 | |||
46 | def create_db(self): |
|
47 | def create_db(self): | |
47 | return DictDB() |
|
48 | return DictDB() | |
48 |
|
49 | |||
49 | def load_records(self, n=1): |
|
50 | def load_records(self, n=1): | |
50 | """load n records for testing""" |
|
51 | """load n records for testing""" | |
51 | #sleep 1/10 s, to ensure timestamp is different to previous calls |
|
52 | #sleep 1/10 s, to ensure timestamp is different to previous calls | |
52 | time.sleep(0.1) |
|
53 | time.sleep(0.1) | |
53 | msg_ids = [] |
|
54 | msg_ids = [] | |
54 | for i in range(n): |
|
55 | for i in range(n): | |
55 | msg = self.session.msg('apply_request', content=dict(a=5)) |
|
56 | msg = self.session.msg('apply_request', content=dict(a=5)) | |
56 | msg['buffers'] = [] |
|
57 | msg['buffers'] = [] | |
57 | rec = init_record(msg) |
|
58 | rec = init_record(msg) | |
58 | msg_id = msg['header']['msg_id'] |
|
59 | msg_id = msg['header']['msg_id'] | |
59 | msg_ids.append(msg_id) |
|
60 | msg_ids.append(msg_id) | |
60 | self.db.add_record(msg_id, rec) |
|
61 | self.db.add_record(msg_id, rec) | |
61 | return msg_ids |
|
62 | return msg_ids | |
62 |
|
63 | |||
63 | def test_add_record(self): |
|
64 | def test_add_record(self): | |
64 | before = self.db.get_history() |
|
65 | before = self.db.get_history() | |
65 | self.load_records(5) |
|
66 | self.load_records(5) | |
66 | after = self.db.get_history() |
|
67 | after = self.db.get_history() | |
67 | self.assertEquals(len(after), len(before)+5) |
|
68 | self.assertEquals(len(after), len(before)+5) | |
68 | self.assertEquals(after[:-5],before) |
|
69 | self.assertEquals(after[:-5],before) | |
69 |
|
70 | |||
70 | def test_drop_record(self): |
|
71 | def test_drop_record(self): | |
71 | msg_id = self.load_records()[-1] |
|
72 | msg_id = self.load_records()[-1] | |
72 | rec = self.db.get_record(msg_id) |
|
73 | rec = self.db.get_record(msg_id) | |
73 | self.db.drop_record(msg_id) |
|
74 | self.db.drop_record(msg_id) | |
74 | self.assertRaises(KeyError,self.db.get_record, msg_id) |
|
75 | self.assertRaises(KeyError,self.db.get_record, msg_id) | |
75 |
|
76 | |||
76 | def _round_to_millisecond(self, dt): |
|
77 | def _round_to_millisecond(self, dt): | |
77 | """necessary because mongodb rounds microseconds""" |
|
78 | """necessary because mongodb rounds microseconds""" | |
78 | micro = dt.microsecond |
|
79 | micro = dt.microsecond | |
79 | extra = int(str(micro)[-3:]) |
|
80 | extra = int(str(micro)[-3:]) | |
80 | return dt - timedelta(microseconds=extra) |
|
81 | return dt - timedelta(microseconds=extra) | |
81 |
|
82 | |||
82 | def test_update_record(self): |
|
83 | def test_update_record(self): | |
83 | now = self._round_to_millisecond(datetime.now()) |
|
84 | now = self._round_to_millisecond(datetime.now()) | |
84 | # |
|
85 | # | |
85 | msg_id = self.db.get_history()[-1] |
|
86 | msg_id = self.db.get_history()[-1] | |
86 | rec1 = self.db.get_record(msg_id) |
|
87 | rec1 = self.db.get_record(msg_id) | |
87 | data = {'stdout': 'hello there', 'completed' : now} |
|
88 | data = {'stdout': 'hello there', 'completed' : now} | |
88 | self.db.update_record(msg_id, data) |
|
89 | self.db.update_record(msg_id, data) | |
89 | rec2 = self.db.get_record(msg_id) |
|
90 | rec2 = self.db.get_record(msg_id) | |
90 | self.assertEquals(rec2['stdout'], 'hello there') |
|
91 | self.assertEquals(rec2['stdout'], 'hello there') | |
91 | self.assertEquals(rec2['completed'], now) |
|
92 | self.assertEquals(rec2['completed'], now) | |
92 | rec1.update(data) |
|
93 | rec1.update(data) | |
93 | self.assertEquals(rec1, rec2) |
|
94 | self.assertEquals(rec1, rec2) | |
94 |
|
95 | |||
95 | # def test_update_record_bad(self): |
|
96 | # def test_update_record_bad(self): | |
96 | # """test updating nonexistant records""" |
|
97 | # """test updating nonexistant records""" | |
97 | # msg_id = str(uuid.uuid4()) |
|
98 | # msg_id = str(uuid.uuid4()) | |
98 | # data = {'stdout': 'hello there'} |
|
99 | # data = {'stdout': 'hello there'} | |
99 | # self.assertRaises(KeyError, self.db.update_record, msg_id, data) |
|
100 | # self.assertRaises(KeyError, self.db.update_record, msg_id, data) | |
100 |
|
101 | |||
101 | def test_find_records_dt(self): |
|
102 | def test_find_records_dt(self): | |
102 | """test finding records by date""" |
|
103 | """test finding records by date""" | |
103 | hist = self.db.get_history() |
|
104 | hist = self.db.get_history() | |
104 | middle = self.db.get_record(hist[len(hist)//2]) |
|
105 | middle = self.db.get_record(hist[len(hist)//2]) | |
105 | tic = middle['submitted'] |
|
106 | tic = middle['submitted'] | |
106 | before = self.db.find_records({'submitted' : {'$lt' : tic}}) |
|
107 | before = self.db.find_records({'submitted' : {'$lt' : tic}}) | |
107 | after = self.db.find_records({'submitted' : {'$gte' : tic}}) |
|
108 | after = self.db.find_records({'submitted' : {'$gte' : tic}}) | |
108 | self.assertEquals(len(before)+len(after),len(hist)) |
|
109 | self.assertEquals(len(before)+len(after),len(hist)) | |
109 | for b in before: |
|
110 | for b in before: | |
110 | self.assertTrue(b['submitted'] < tic) |
|
111 | self.assertTrue(b['submitted'] < tic) | |
111 | for a in after: |
|
112 | for a in after: | |
112 | self.assertTrue(a['submitted'] >= tic) |
|
113 | self.assertTrue(a['submitted'] >= tic) | |
113 | same = self.db.find_records({'submitted' : tic}) |
|
114 | same = self.db.find_records({'submitted' : tic}) | |
114 | for s in same: |
|
115 | for s in same: | |
115 | self.assertTrue(s['submitted'] == tic) |
|
116 | self.assertTrue(s['submitted'] == tic) | |
116 |
|
117 | |||
117 | def test_find_records_keys(self): |
|
118 | def test_find_records_keys(self): | |
118 | """test extracting subset of record keys""" |
|
119 | """test extracting subset of record keys""" | |
119 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted', 'completed']) |
|
120 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted', 'completed']) | |
120 | for rec in found: |
|
121 | for rec in found: | |
121 | self.assertEquals(set(rec.keys()), set(['msg_id', 'submitted', 'completed'])) |
|
122 | self.assertEquals(set(rec.keys()), set(['msg_id', 'submitted', 'completed'])) | |
122 |
|
123 | |||
123 | def test_find_records_msg_id(self): |
|
124 | def test_find_records_msg_id(self): | |
124 | """ensure msg_id is always in found records""" |
|
125 | """ensure msg_id is always in found records""" | |
125 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted', 'completed']) |
|
126 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted', 'completed']) | |
126 | for rec in found: |
|
127 | for rec in found: | |
127 | self.assertTrue('msg_id' in rec.keys()) |
|
128 | self.assertTrue('msg_id' in rec.keys()) | |
128 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted']) |
|
129 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['submitted']) | |
129 | for rec in found: |
|
130 | for rec in found: | |
130 | self.assertTrue('msg_id' in rec.keys()) |
|
131 | self.assertTrue('msg_id' in rec.keys()) | |
131 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['msg_id']) |
|
132 | found = self.db.find_records({'msg_id': {'$ne' : ''}},keys=['msg_id']) | |
132 | for rec in found: |
|
133 | for rec in found: | |
133 | self.assertTrue('msg_id' in rec.keys()) |
|
134 | self.assertTrue('msg_id' in rec.keys()) | |
134 |
|
135 | |||
135 | def test_find_records_in(self): |
|
136 | def test_find_records_in(self): | |
136 | """test finding records with '$in','$nin' operators""" |
|
137 | """test finding records with '$in','$nin' operators""" | |
137 | hist = self.db.get_history() |
|
138 | hist = self.db.get_history() | |
138 | even = hist[::2] |
|
139 | even = hist[::2] | |
139 | odd = hist[1::2] |
|
140 | odd = hist[1::2] | |
140 | recs = self.db.find_records({ 'msg_id' : {'$in' : even}}) |
|
141 | recs = self.db.find_records({ 'msg_id' : {'$in' : even}}) | |
141 | found = [ r['msg_id'] for r in recs ] |
|
142 | found = [ r['msg_id'] for r in recs ] | |
142 | self.assertEquals(set(even), set(found)) |
|
143 | self.assertEquals(set(even), set(found)) | |
143 | recs = self.db.find_records({ 'msg_id' : {'$nin' : even}}) |
|
144 | recs = self.db.find_records({ 'msg_id' : {'$nin' : even}}) | |
144 | found = [ r['msg_id'] for r in recs ] |
|
145 | found = [ r['msg_id'] for r in recs ] | |
145 | self.assertEquals(set(odd), set(found)) |
|
146 | self.assertEquals(set(odd), set(found)) | |
146 |
|
147 | |||
147 | def test_get_history(self): |
|
148 | def test_get_history(self): | |
148 | msg_ids = self.db.get_history() |
|
149 | msg_ids = self.db.get_history() | |
149 | latest = datetime(1984,1,1) |
|
150 | latest = datetime(1984,1,1) | |
150 | for msg_id in msg_ids: |
|
151 | for msg_id in msg_ids: | |
151 | rec = self.db.get_record(msg_id) |
|
152 | rec = self.db.get_record(msg_id) | |
152 | newt = rec['submitted'] |
|
153 | newt = rec['submitted'] | |
153 | self.assertTrue(newt >= latest) |
|
154 | self.assertTrue(newt >= latest) | |
154 | latest = newt |
|
155 | latest = newt | |
155 | msg_id = self.load_records(1)[-1] |
|
156 | msg_id = self.load_records(1)[-1] | |
156 | self.assertEquals(self.db.get_history()[-1],msg_id) |
|
157 | self.assertEquals(self.db.get_history()[-1],msg_id) | |
157 |
|
158 | |||
158 | def test_datetime(self): |
|
159 | def test_datetime(self): | |
159 | """get/set timestamps with datetime objects""" |
|
160 | """get/set timestamps with datetime objects""" | |
160 | msg_id = self.db.get_history()[-1] |
|
161 | msg_id = self.db.get_history()[-1] | |
161 | rec = self.db.get_record(msg_id) |
|
162 | rec = self.db.get_record(msg_id) | |
162 | self.assertTrue(isinstance(rec['submitted'], datetime)) |
|
163 | self.assertTrue(isinstance(rec['submitted'], datetime)) | |
163 | self.db.update_record(msg_id, dict(completed=datetime.now())) |
|
164 | self.db.update_record(msg_id, dict(completed=datetime.now())) | |
164 | rec = self.db.get_record(msg_id) |
|
165 | rec = self.db.get_record(msg_id) | |
165 | self.assertTrue(isinstance(rec['completed'], datetime)) |
|
166 | self.assertTrue(isinstance(rec['completed'], datetime)) | |
166 |
|
167 | |||
167 | def test_drop_matching(self): |
|
168 | def test_drop_matching(self): | |
168 | msg_ids = self.load_records(10) |
|
169 | msg_ids = self.load_records(10) | |
169 | query = {'msg_id' : {'$in':msg_ids}} |
|
170 | query = {'msg_id' : {'$in':msg_ids}} | |
170 | self.db.drop_matching_records(query) |
|
171 | self.db.drop_matching_records(query) | |
171 | recs = self.db.find_records(query) |
|
172 | recs = self.db.find_records(query) | |
172 | self.assertEquals(len(recs), 0) |
|
173 | self.assertEquals(len(recs), 0) | |
173 |
|
174 | |||
174 | def test_null(self): |
|
175 | def test_null(self): | |
175 | """test None comparison queries""" |
|
176 | """test None comparison queries""" | |
176 | msg_ids = self.load_records(10) |
|
177 | msg_ids = self.load_records(10) | |
177 |
|
178 | |||
178 | query = {'msg_id' : None} |
|
179 | query = {'msg_id' : None} | |
179 | recs = self.db.find_records(query) |
|
180 | recs = self.db.find_records(query) | |
180 | self.assertEquals(len(recs), 0) |
|
181 | self.assertEquals(len(recs), 0) | |
181 |
|
182 | |||
182 | query = {'msg_id' : {'$ne' : None}} |
|
183 | query = {'msg_id' : {'$ne' : None}} | |
183 | recs = self.db.find_records(query) |
|
184 | recs = self.db.find_records(query) | |
184 | self.assertTrue(len(recs) >= 10) |
|
185 | self.assertTrue(len(recs) >= 10) | |
185 |
|
186 | |||
186 | def test_pop_safe_get(self): |
|
187 | def test_pop_safe_get(self): | |
187 | """editing query results shouldn't affect record [get]""" |
|
188 | """editing query results shouldn't affect record [get]""" | |
188 | msg_id = self.db.get_history()[-1] |
|
189 | msg_id = self.db.get_history()[-1] | |
189 | rec = self.db.get_record(msg_id) |
|
190 | rec = self.db.get_record(msg_id) | |
190 | rec.pop('buffers') |
|
191 | rec.pop('buffers') | |
191 | rec['garbage'] = 'hello' |
|
192 | rec['garbage'] = 'hello' | |
192 | rec2 = self.db.get_record(msg_id) |
|
193 | rec2 = self.db.get_record(msg_id) | |
193 | self.assertTrue('buffers' in rec2) |
|
194 | self.assertTrue('buffers' in rec2) | |
194 | self.assertFalse('garbage' in rec2) |
|
195 | self.assertFalse('garbage' in rec2) | |
195 |
|
196 | |||
196 | def test_pop_safe_find(self): |
|
197 | def test_pop_safe_find(self): | |
197 | """editing query results shouldn't affect record [find]""" |
|
198 | """editing query results shouldn't affect record [find]""" | |
198 | msg_id = self.db.get_history()[-1] |
|
199 | msg_id = self.db.get_history()[-1] | |
199 | rec = self.db.find_records({'msg_id' : msg_id})[0] |
|
200 | rec = self.db.find_records({'msg_id' : msg_id})[0] | |
200 | rec.pop('buffers') |
|
201 | rec.pop('buffers') | |
201 | rec['garbage'] = 'hello' |
|
202 | rec['garbage'] = 'hello' | |
202 | rec2 = self.db.find_records({'msg_id' : msg_id})[0] |
|
203 | rec2 = self.db.find_records({'msg_id' : msg_id})[0] | |
203 | self.assertTrue('buffers' in rec2) |
|
204 | self.assertTrue('buffers' in rec2) | |
204 | self.assertFalse('garbage' in rec2) |
|
205 | self.assertFalse('garbage' in rec2) | |
205 |
|
206 | |||
206 | def test_pop_safe_find_keys(self): |
|
207 | def test_pop_safe_find_keys(self): | |
207 | """editing query results shouldn't affect record [find+keys]""" |
|
208 | """editing query results shouldn't affect record [find+keys]""" | |
208 | msg_id = self.db.get_history()[-1] |
|
209 | msg_id = self.db.get_history()[-1] | |
209 | rec = self.db.find_records({'msg_id' : msg_id}, keys=['buffers'])[0] |
|
210 | rec = self.db.find_records({'msg_id' : msg_id}, keys=['buffers'])[0] | |
210 | rec.pop('buffers') |
|
211 | rec.pop('buffers') | |
211 | rec['garbage'] = 'hello' |
|
212 | rec['garbage'] = 'hello' | |
212 | rec2 = self.db.find_records({'msg_id' : msg_id})[0] |
|
213 | rec2 = self.db.find_records({'msg_id' : msg_id})[0] | |
213 | self.assertTrue('buffers' in rec2) |
|
214 | self.assertTrue('buffers' in rec2) | |
214 | self.assertFalse('garbage' in rec2) |
|
215 | self.assertFalse('garbage' in rec2) | |
215 |
|
216 | |||
216 |
|
217 | |||
217 | class TestSQLiteBackend(TestDictBackend): |
|
218 | class TestSQLiteBackend(TestDictBackend): | |
218 |
|
219 | |||
219 | @dec.skip_without('sqlite3') |
|
220 | @dec.skip_without('sqlite3') | |
220 | def create_db(self): |
|
221 | def create_db(self): | |
221 | return SQLiteDB(location=tempfile.gettempdir()) |
|
222 | return SQLiteDB(location=tempfile.gettempdir()) | |
222 |
|
223 | |||
223 | def tearDown(self): |
|
224 | def tearDown(self): | |
224 | self.db._db.close() |
|
225 | self.db._db.close() | |
|
226 | ||||
|
227 | ||||
|
228 | def teardown(): | |||
|
229 | """cleanup task db file after all tests have run""" | |||
|
230 | try: | |||
|
231 | os.remove(os.path.join(tempfile.gettempdir(), 'tasks.db')) | |||
|
232 | except: | |||
|
233 | pass |
General Comments 0
You need to be logged in to leave comments.
Login now