/brz/remove-bazaar

To get this branch, use:
bzr branch http://gegoxaren.bato24.eu/bzr/brz/remove-bazaar

« back to all changes in this revision

Viewing changes to bzrlib/tests/test_btree_index.py

  • Committer: John Arbash Meinel
  • Date: 2009-12-22 16:28:47 UTC
  • mto: This revision was merged to the branch mainline in revision 4922.
  • Revision ID: john@arbash-meinel.com-20091222162847-tvnsc69to4l4uf5r
Implement a permute_for_extension helper.

Use it for all of the 'simple' extension permutations.
It basically permutes all tests in the current module, by setting TestCase.module.
Which works well for most of our extension tests. Some had more advanced
handling of permutations (extra permutations, custom vars, etc.)

Show diffs side-by-side

added added

removed removed

Lines of Context:
 
1
# Copyright (C) 2008, 2009 Canonical Ltd
 
2
#
 
3
# This program is free software; you can redistribute it and/or modify
 
4
# it under the terms of the GNU General Public License as published by
 
5
# the Free Software Foundation; either version 2 of the License, or
 
6
# (at your option) any later version.
 
7
#
 
8
# This program is distributed in the hope that it will be useful,
 
9
# but WITHOUT ANY WARRANTY; without even the implied warranty of
 
10
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 
11
# GNU General Public License for more details.
 
12
#
 
13
# You should have received a copy of the GNU General Public License
 
14
# along with this program; if not, write to the Free Software
 
15
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 
16
#
 
17
 
 
18
"""Tests for btree indices."""
 
19
 
 
20
import pprint
 
21
import zlib
 
22
 
 
23
from bzrlib import (
 
24
    btree_index,
 
25
    errors,
 
26
    fifo_cache,
 
27
    lru_cache,
 
28
    osutils,
 
29
    tests,
 
30
    )
 
31
from bzrlib.tests import (
 
32
    TestCaseWithTransport,
 
33
    condition_isinstance,
 
34
    multiply_tests,
 
35
    split_suite_by_condition,
 
36
    )
 
37
from bzrlib.transport import get_transport
 
38
 
 
39
 
 
40
def load_tests(standard_tests, module, loader):
 
41
    # parameterise the TestBTreeNodes tests
 
42
    node_tests, others = split_suite_by_condition(standard_tests,
 
43
        condition_isinstance(TestBTreeNodes))
 
44
    import bzrlib._btree_serializer_py as py_module
 
45
    scenarios = [('python', {'parse_btree': py_module})]
 
46
    if compiled_btreeparser_feature.available():
 
47
        scenarios.append(('C', {'parse_btree':
 
48
                                compiled_btreeparser_feature.module}))
 
49
    return multiply_tests(node_tests, scenarios, others)
 
50
 
 
51
 
 
52
compiled_btreeparser_feature = tests.ModuleAvailableFeature(
 
53
                                'bzrlib._btree_serializer_pyx')
 
54
 
 
55
 
 
56
class BTreeTestCase(TestCaseWithTransport):
 
57
    # test names here are suffixed by the key length and reference list count
 
58
    # that they test.
 
59
 
 
60
    def setUp(self):
 
61
        TestCaseWithTransport.setUp(self)
 
62
        self._original_header = btree_index._RESERVED_HEADER_BYTES
 
63
        def restore():
 
64
            btree_index._RESERVED_HEADER_BYTES = self._original_header
 
65
        self.addCleanup(restore)
 
66
        btree_index._RESERVED_HEADER_BYTES = 100
 
67
 
 
68
    def make_nodes(self, count, key_elements, reference_lists):
 
69
        """Generate count*key_elements sample nodes."""
 
70
        keys = []
 
71
        for prefix_pos in range(key_elements):
 
72
            if key_elements - 1:
 
73
                prefix = (str(prefix_pos) * 40,)
 
74
            else:
 
75
                prefix = ()
 
76
            for pos in xrange(count):
 
77
                # TODO: This creates odd keys. When count == 100,000, it
 
78
                #       creates a 240 byte key
 
79
                key = prefix + (str(pos) * 40,)
 
80
                value = "value:%s" % pos
 
81
                if reference_lists:
 
82
                    # generate some references
 
83
                    refs = []
 
84
                    for list_pos in range(reference_lists):
 
85
                        # as many keys in each list as its index + the key depth
 
86
                        # mod 2 - this generates both 0 length lists and
 
87
                        # ones slightly longer than the number of lists.
 
88
                        # It also ensures we have non homogeneous lists.
 
89
                        refs.append([])
 
90
                        for ref_pos in range(list_pos + pos % 2):
 
91
                            if pos % 2:
 
92
                                # refer to a nearby key
 
93
                                refs[-1].append(prefix + ("ref" + str(pos - 1) * 40,))
 
94
                            else:
 
95
                                # serial of this ref in the ref list
 
96
                                refs[-1].append(prefix + ("ref" + str(ref_pos) * 40,))
 
97
                        refs[-1] = tuple(refs[-1])
 
98
                    refs = tuple(refs)
 
99
                else:
 
100
                    refs = ()
 
101
                keys.append((key, value, refs))
 
102
        return keys
 
103
 
 
104
    def shrink_page_size(self):
 
105
        """Shrink the default page size so that less fits in a page."""
 
106
        old_page_size = btree_index._PAGE_SIZE
 
107
        def cleanup():
 
108
            btree_index._PAGE_SIZE = old_page_size
 
109
        self.addCleanup(cleanup)
 
110
        btree_index._PAGE_SIZE = 2048
 
111
 
 
112
 
 
113
class TestBTreeBuilder(BTreeTestCase):
 
114
 
 
115
    def test_clear_cache(self):
 
116
        builder = btree_index.BTreeBuilder(reference_lists=0, key_elements=1)
 
117
        # This is a no-op, but we need the api to be consistent with other
 
118
        # BTreeGraphIndex apis.
 
119
        builder.clear_cache()
 
120
 
 
121
    def test_empty_1_0(self):
 
122
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
 
123
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
124
        temp_file = builder.finish()
 
125
        content = temp_file.read()
 
126
        del temp_file
 
127
        self.assertEqual(
 
128
            "B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=0\n"
 
129
            "row_lengths=\n",
 
130
            content)
 
131
 
 
132
    def test_empty_2_1(self):
 
133
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=1)
 
134
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
135
        temp_file = builder.finish()
 
136
        content = temp_file.read()
 
137
        del temp_file
 
138
        self.assertEqual(
 
139
            "B+Tree Graph Index 2\nnode_ref_lists=1\nkey_elements=2\nlen=0\n"
 
140
            "row_lengths=\n",
 
141
            content)
 
142
 
 
143
    def test_root_leaf_1_0(self):
 
144
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
 
145
        nodes = self.make_nodes(5, 1, 0)
 
146
        for node in nodes:
 
147
            builder.add_node(*node)
 
148
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
149
        temp_file = builder.finish()
 
150
        content = temp_file.read()
 
151
        del temp_file
 
152
        self.assertEqual(131, len(content))
 
153
        self.assertEqual(
 
154
            "B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=5\n"
 
155
            "row_lengths=1\n",
 
156
            content[:73])
 
157
        node_content = content[73:]
 
158
        node_bytes = zlib.decompress(node_content)
 
159
        expected_node = ("type=leaf\n"
 
160
            "0000000000000000000000000000000000000000\x00\x00value:0\n"
 
161
            "1111111111111111111111111111111111111111\x00\x00value:1\n"
 
162
            "2222222222222222222222222222222222222222\x00\x00value:2\n"
 
163
            "3333333333333333333333333333333333333333\x00\x00value:3\n"
 
164
            "4444444444444444444444444444444444444444\x00\x00value:4\n")
 
165
        self.assertEqual(expected_node, node_bytes)
 
166
 
 
167
    def test_root_leaf_2_2(self):
 
168
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
169
        nodes = self.make_nodes(5, 2, 2)
 
170
        for node in nodes:
 
171
            builder.add_node(*node)
 
172
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
173
        temp_file = builder.finish()
 
174
        content = temp_file.read()
 
175
        del temp_file
 
176
        self.assertEqual(238, len(content))
 
177
        self.assertEqual(
 
178
            "B+Tree Graph Index 2\nnode_ref_lists=2\nkey_elements=2\nlen=10\n"
 
179
            "row_lengths=1\n",
 
180
            content[:74])
 
181
        node_content = content[74:]
 
182
        node_bytes = zlib.decompress(node_content)
 
183
        expected_node = (
 
184
            "type=leaf\n"
 
185
            "0000000000000000000000000000000000000000\x000000000000000000000000000000000000000000\x00\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:0\n"
 
186
            "0000000000000000000000000000000000000000\x001111111111111111111111111111111111111111\x000000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\r0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:1\n"
 
187
            "0000000000000000000000000000000000000000\x002222222222222222222222222222222222222222\x00\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:2\n"
 
188
            "0000000000000000000000000000000000000000\x003333333333333333333333333333333333333333\x000000000000000000000000000000000000000000\x00ref2222222222222222222222222222222222222222\t0000000000000000000000000000000000000000\x00ref2222222222222222222222222222222222222222\r0000000000000000000000000000000000000000\x00ref2222222222222222222222222222222222222222\x00value:3\n"
 
189
            "0000000000000000000000000000000000000000\x004444444444444444444444444444444444444444\x00\t0000000000000000000000000000000000000000\x00ref0000000000000000000000000000000000000000\x00value:4\n"
 
190
            "1111111111111111111111111111111111111111\x000000000000000000000000000000000000000000\x00\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:0\n"
 
191
            "1111111111111111111111111111111111111111\x001111111111111111111111111111111111111111\x001111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\r1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:1\n"
 
192
            "1111111111111111111111111111111111111111\x002222222222222222222222222222222222222222\x00\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:2\n"
 
193
            "1111111111111111111111111111111111111111\x003333333333333333333333333333333333333333\x001111111111111111111111111111111111111111\x00ref2222222222222222222222222222222222222222\t1111111111111111111111111111111111111111\x00ref2222222222222222222222222222222222222222\r1111111111111111111111111111111111111111\x00ref2222222222222222222222222222222222222222\x00value:3\n"
 
194
            "1111111111111111111111111111111111111111\x004444444444444444444444444444444444444444\x00\t1111111111111111111111111111111111111111\x00ref0000000000000000000000000000000000000000\x00value:4\n"
 
195
            ""
 
196
            )
 
197
        self.assertEqual(expected_node, node_bytes)
 
198
 
 
199
    def test_2_leaves_1_0(self):
 
200
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
 
201
        nodes = self.make_nodes(400, 1, 0)
 
202
        for node in nodes:
 
203
            builder.add_node(*node)
 
204
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
205
        temp_file = builder.finish()
 
206
        content = temp_file.read()
 
207
        del temp_file
 
208
        self.assertEqual(9283, len(content))
 
209
        self.assertEqual(
 
210
            "B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=400\n"
 
211
            "row_lengths=1,2\n",
 
212
            content[:77])
 
213
        root = content[77:4096]
 
214
        leaf1 = content[4096:8192]
 
215
        leaf2 = content[8192:]
 
216
        root_bytes = zlib.decompress(root)
 
217
        expected_root = (
 
218
            "type=internal\n"
 
219
            "offset=0\n"
 
220
            ) + ("307" * 40) + "\n"
 
221
        self.assertEqual(expected_root, root_bytes)
 
222
        # We already know serialisation works for leaves, check key selection:
 
223
        leaf1_bytes = zlib.decompress(leaf1)
 
224
        sorted_node_keys = sorted(node[0] for node in nodes)
 
225
        node = btree_index._LeafNode(leaf1_bytes, 1, 0)
 
226
        self.assertEqual(231, len(node.keys))
 
227
        self.assertEqual(sorted_node_keys[:231], sorted(node.keys))
 
228
        leaf2_bytes = zlib.decompress(leaf2)
 
229
        node = btree_index._LeafNode(leaf2_bytes, 1, 0)
 
230
        self.assertEqual(400 - 231, len(node.keys))
 
231
        self.assertEqual(sorted_node_keys[231:], sorted(node.keys))
 
232
 
 
233
    def test_last_page_rounded_1_layer(self):
 
234
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
 
235
        nodes = self.make_nodes(10, 1, 0)
 
236
        for node in nodes:
 
237
            builder.add_node(*node)
 
238
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
239
        temp_file = builder.finish()
 
240
        content = temp_file.read()
 
241
        del temp_file
 
242
        self.assertEqual(155, len(content))
 
243
        self.assertEqual(
 
244
            "B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=10\n"
 
245
            "row_lengths=1\n",
 
246
            content[:74])
 
247
        # Check thelast page is well formed
 
248
        leaf2 = content[74:]
 
249
        leaf2_bytes = zlib.decompress(leaf2)
 
250
        node = btree_index._LeafNode(leaf2_bytes, 1, 0)
 
251
        self.assertEqual(10, len(node.keys))
 
252
        sorted_node_keys = sorted(node[0] for node in nodes)
 
253
        self.assertEqual(sorted_node_keys, sorted(node.keys))
 
254
 
 
255
    def test_last_page_not_rounded_2_layer(self):
 
256
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
 
257
        nodes = self.make_nodes(400, 1, 0)
 
258
        for node in nodes:
 
259
            builder.add_node(*node)
 
260
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
261
        temp_file = builder.finish()
 
262
        content = temp_file.read()
 
263
        del temp_file
 
264
        self.assertEqual(9283, len(content))
 
265
        self.assertEqual(
 
266
            "B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=400\n"
 
267
            "row_lengths=1,2\n",
 
268
            content[:77])
 
269
        # Check the last page is well formed
 
270
        leaf2 = content[8192:]
 
271
        leaf2_bytes = zlib.decompress(leaf2)
 
272
        node = btree_index._LeafNode(leaf2_bytes, 1, 0)
 
273
        self.assertEqual(400 - 231, len(node.keys))
 
274
        sorted_node_keys = sorted(node[0] for node in nodes)
 
275
        self.assertEqual(sorted_node_keys[231:], sorted(node.keys))
 
276
 
 
277
    def test_three_level_tree_details(self):
 
278
        # The left most pointer in the second internal node in a row should
 
279
        # pointer to the second node that the internal node is for, _not_
 
280
        # the first, otherwise the first node overlaps with the last node of
 
281
        # the prior internal node on that row.
 
282
        self.shrink_page_size()
 
283
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
284
        # 40K nodes is enough to create a two internal nodes on the second
 
285
        # level, with a 2K page size
 
286
        nodes = self.make_nodes(20000, 2, 2)
 
287
 
 
288
        for node in nodes:
 
289
            builder.add_node(*node)
 
290
        transport = get_transport('trace+' + self.get_url(''))
 
291
        size = transport.put_file('index', self.time(builder.finish))
 
292
        del builder
 
293
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
294
        # Seed the metadata, we're using internal calls now.
 
295
        index.key_count()
 
296
        self.assertEqual(3, len(index._row_lengths),
 
297
            "Not enough rows: %r" % index._row_lengths)
 
298
        self.assertEqual(4, len(index._row_offsets))
 
299
        self.assertEqual(sum(index._row_lengths), index._row_offsets[-1])
 
300
        internal_nodes = index._get_internal_nodes([0, 1, 2])
 
301
        root_node = internal_nodes[0]
 
302
        internal_node1 = internal_nodes[1]
 
303
        internal_node2 = internal_nodes[2]
 
304
        # The left most node node2 points at should be one after the right most
 
305
        # node pointed at by node1.
 
306
        self.assertEqual(internal_node2.offset, 1 + len(internal_node1.keys))
 
307
        # The left most key of the second node pointed at by internal_node2
 
308
        # should be its first key. We can check this by looking for its first key
 
309
        # in the second node it points at
 
310
        pos = index._row_offsets[2] + internal_node2.offset + 1
 
311
        leaf = index._get_leaf_nodes([pos])[pos]
 
312
        self.assertTrue(internal_node2.keys[0] in leaf.keys)
 
313
 
 
314
    def test_2_leaves_2_2(self):
 
315
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
316
        nodes = self.make_nodes(100, 2, 2)
 
317
        for node in nodes:
 
318
            builder.add_node(*node)
 
319
        # NamedTemporaryFile dies on builder.finish().read(). weird.
 
320
        temp_file = builder.finish()
 
321
        content = temp_file.read()
 
322
        del temp_file
 
323
        self.assertEqual(12643, len(content))
 
324
        self.assertEqual(
 
325
            "B+Tree Graph Index 2\nnode_ref_lists=2\nkey_elements=2\nlen=200\n"
 
326
            "row_lengths=1,3\n",
 
327
            content[:77])
 
328
        root = content[77:4096]
 
329
        leaf1 = content[4096:8192]
 
330
        leaf2 = content[8192:12288]
 
331
        leaf3 = content[12288:]
 
332
        root_bytes = zlib.decompress(root)
 
333
        expected_root = (
 
334
            "type=internal\n"
 
335
            "offset=0\n"
 
336
            + ("0" * 40) + "\x00" + ("91" * 40) + "\n"
 
337
            + ("1" * 40) + "\x00" + ("81" * 40) + "\n"
 
338
            )
 
339
        self.assertEqual(expected_root, root_bytes)
 
340
        # We assume the other leaf nodes have been written correctly - layering
 
341
        # FTW.
 
342
 
 
343
    def test_spill_index_stress_1_1(self):
 
344
        builder = btree_index.BTreeBuilder(key_elements=1, spill_at=2)
 
345
        nodes = [node[0:2] for node in self.make_nodes(16, 1, 0)]
 
346
        builder.add_node(*nodes[0])
 
347
        # Test the parts of the index that take up memory are doing so
 
348
        # predictably.
 
349
        self.assertEqual(1, len(builder._nodes))
 
350
        self.assertIs(None, builder._nodes_by_key)
 
351
        builder.add_node(*nodes[1])
 
352
        self.assertEqual(0, len(builder._nodes))
 
353
        self.assertIs(None, builder._nodes_by_key)
 
354
        self.assertEqual(1, len(builder._backing_indices))
 
355
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
356
        # now back to memory
 
357
        builder.add_node(*nodes[2])
 
358
        self.assertEqual(1, len(builder._nodes))
 
359
        self.assertIs(None, builder._nodes_by_key)
 
360
        # And spills to a second backing index combing all
 
361
        builder.add_node(*nodes[3])
 
362
        self.assertEqual(0, len(builder._nodes))
 
363
        self.assertIs(None, builder._nodes_by_key)
 
364
        self.assertEqual(2, len(builder._backing_indices))
 
365
        self.assertEqual(None, builder._backing_indices[0])
 
366
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
367
        # The next spills to the 2-len slot
 
368
        builder.add_node(*nodes[4])
 
369
        builder.add_node(*nodes[5])
 
370
        self.assertEqual(0, len(builder._nodes))
 
371
        self.assertIs(None, builder._nodes_by_key)
 
372
        self.assertEqual(2, len(builder._backing_indices))
 
373
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
374
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
375
        # Next spill combines
 
376
        builder.add_node(*nodes[6])
 
377
        builder.add_node(*nodes[7])
 
378
        self.assertEqual(3, len(builder._backing_indices))
 
379
        self.assertEqual(None, builder._backing_indices[0])
 
380
        self.assertEqual(None, builder._backing_indices[1])
 
381
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
382
        # And so forth - counting up in binary.
 
383
        builder.add_node(*nodes[8])
 
384
        builder.add_node(*nodes[9])
 
385
        self.assertEqual(3, len(builder._backing_indices))
 
386
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
387
        self.assertEqual(None, builder._backing_indices[1])
 
388
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
389
        builder.add_node(*nodes[10])
 
390
        builder.add_node(*nodes[11])
 
391
        self.assertEqual(3, len(builder._backing_indices))
 
392
        self.assertEqual(None, builder._backing_indices[0])
 
393
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
394
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
395
        builder.add_node(*nodes[12])
 
396
        # Test that memory and disk are both used for query methods; and that
 
397
        # None is skipped over happily.
 
398
        self.assertEqual([(builder,) + node for node in sorted(nodes[:13])],
 
399
            list(builder.iter_all_entries()))
 
400
        # Two nodes - one memory one disk
 
401
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
 
402
            set(builder.iter_entries([nodes[12][0], nodes[11][0]])))
 
403
        self.assertEqual(13, builder.key_count())
 
404
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
 
405
            set(builder.iter_entries_prefix([nodes[12][0], nodes[11][0]])))
 
406
        builder.add_node(*nodes[13])
 
407
        self.assertEqual(3, len(builder._backing_indices))
 
408
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
409
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
410
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
411
        builder.add_node(*nodes[14])
 
412
        builder.add_node(*nodes[15])
 
413
        self.assertEqual(4, len(builder._backing_indices))
 
414
        self.assertEqual(None, builder._backing_indices[0])
 
415
        self.assertEqual(None, builder._backing_indices[1])
 
416
        self.assertEqual(None, builder._backing_indices[2])
 
417
        self.assertEqual(16, builder._backing_indices[3].key_count())
 
418
        # Now finish, and check we got a correctly ordered tree
 
419
        transport = self.get_transport('')
 
420
        size = transport.put_file('index', builder.finish())
 
421
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
422
        nodes = list(index.iter_all_entries())
 
423
        self.assertEqual(sorted(nodes), nodes)
 
424
        self.assertEqual(16, len(nodes))
 
425
 
 
426
    def test_spill_index_stress_1_1_no_combine(self):
 
427
        builder = btree_index.BTreeBuilder(key_elements=1, spill_at=2)
 
428
        builder.set_optimize(for_size=False, combine_backing_indices=False)
 
429
        nodes = [node[0:2] for node in self.make_nodes(16, 1, 0)]
 
430
        builder.add_node(*nodes[0])
 
431
        # Test the parts of the index that take up memory are doing so
 
432
        # predictably.
 
433
        self.assertEqual(1, len(builder._nodes))
 
434
        self.assertIs(None, builder._nodes_by_key)
 
435
        builder.add_node(*nodes[1])
 
436
        self.assertEqual(0, len(builder._nodes))
 
437
        self.assertIs(None, builder._nodes_by_key)
 
438
        self.assertEqual(1, len(builder._backing_indices))
 
439
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
440
        # now back to memory
 
441
        builder.add_node(*nodes[2])
 
442
        self.assertEqual(1, len(builder._nodes))
 
443
        self.assertIs(None, builder._nodes_by_key)
 
444
        # And spills to a second backing index but doesn't combine
 
445
        builder.add_node(*nodes[3])
 
446
        self.assertEqual(0, len(builder._nodes))
 
447
        self.assertIs(None, builder._nodes_by_key)
 
448
        self.assertEqual(2, len(builder._backing_indices))
 
449
        for backing_index in builder._backing_indices:
 
450
            self.assertEqual(2, backing_index.key_count())
 
451
        # The next spills to the 3rd slot
 
452
        builder.add_node(*nodes[4])
 
453
        builder.add_node(*nodes[5])
 
454
        self.assertEqual(0, len(builder._nodes))
 
455
        self.assertIs(None, builder._nodes_by_key)
 
456
        self.assertEqual(3, len(builder._backing_indices))
 
457
        for backing_index in builder._backing_indices:
 
458
            self.assertEqual(2, backing_index.key_count())
 
459
        # Now spill a few more, and check that we don't combine
 
460
        builder.add_node(*nodes[6])
 
461
        builder.add_node(*nodes[7])
 
462
        builder.add_node(*nodes[8])
 
463
        builder.add_node(*nodes[9])
 
464
        builder.add_node(*nodes[10])
 
465
        builder.add_node(*nodes[11])
 
466
        builder.add_node(*nodes[12])
 
467
        self.assertEqual(6, len(builder._backing_indices))
 
468
        for backing_index in builder._backing_indices:
 
469
            self.assertEqual(2, backing_index.key_count())
 
470
        # Test that memory and disk are both used for query methods; and that
 
471
        # None is skipped over happily.
 
472
        self.assertEqual([(builder,) + node for node in sorted(nodes[:13])],
 
473
            list(builder.iter_all_entries()))
 
474
        # Two nodes - one memory one disk
 
475
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
 
476
            set(builder.iter_entries([nodes[12][0], nodes[11][0]])))
 
477
        self.assertEqual(13, builder.key_count())
 
478
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
 
479
            set(builder.iter_entries_prefix([nodes[12][0], nodes[11][0]])))
 
480
        builder.add_node(*nodes[13])
 
481
        builder.add_node(*nodes[14])
 
482
        builder.add_node(*nodes[15])
 
483
        self.assertEqual(8, len(builder._backing_indices))
 
484
        for backing_index in builder._backing_indices:
 
485
            self.assertEqual(2, backing_index.key_count())
 
486
        # Now finish, and check we got a correctly ordered tree
 
487
        transport = self.get_transport('')
 
488
        size = transport.put_file('index', builder.finish())
 
489
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
490
        nodes = list(index.iter_all_entries())
 
491
        self.assertEqual(sorted(nodes), nodes)
 
492
        self.assertEqual(16, len(nodes))
 
493
 
 
494
    def test_set_optimize(self):
 
495
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
496
        builder.set_optimize(for_size=True)
 
497
        self.assertTrue(builder._optimize_for_size)
 
498
        builder.set_optimize(for_size=False)
 
499
        self.assertFalse(builder._optimize_for_size)
 
500
        # test that we can set combine_backing_indices without effecting
 
501
        # _optimize_for_size
 
502
        obj = object()
 
503
        builder._optimize_for_size = obj
 
504
        builder.set_optimize(combine_backing_indices=False)
 
505
        self.assertFalse(builder._combine_backing_indices)
 
506
        self.assertIs(obj, builder._optimize_for_size)
 
507
        builder.set_optimize(combine_backing_indices=True)
 
508
        self.assertTrue(builder._combine_backing_indices)
 
509
        self.assertIs(obj, builder._optimize_for_size)
 
510
 
 
511
    def test_spill_index_stress_2_2(self):
 
512
        # test that references and longer keys don't confuse things.
 
513
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2,
 
514
            spill_at=2)
 
515
        nodes = self.make_nodes(16, 2, 2)
 
516
        builder.add_node(*nodes[0])
 
517
        # Test the parts of the index that take up memory are doing so
 
518
        # predictably.
 
519
        self.assertEqual(1, len(builder._nodes))
 
520
        self.assertIs(None, builder._nodes_by_key)
 
521
        builder.add_node(*nodes[1])
 
522
        self.assertEqual(0, len(builder._nodes))
 
523
        self.assertIs(None, builder._nodes_by_key)
 
524
        self.assertEqual(1, len(builder._backing_indices))
 
525
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
526
        # now back to memory
 
527
        old = dict(builder._get_nodes_by_key()) #Build up the nodes by key dict
 
528
        builder.add_node(*nodes[2])
 
529
        self.assertEqual(1, len(builder._nodes))
 
530
        self.assertIsNot(None, builder._nodes_by_key)
 
531
        self.assertNotEqual({}, builder._nodes_by_key)
 
532
        # We should have a new entry
 
533
        self.assertNotEqual(old, builder._nodes_by_key)
 
534
        # And spills to a second backing index combing all
 
535
        builder.add_node(*nodes[3])
 
536
        self.assertEqual(0, len(builder._nodes))
 
537
        self.assertIs(None, builder._nodes_by_key)
 
538
        self.assertEqual(2, len(builder._backing_indices))
 
539
        self.assertEqual(None, builder._backing_indices[0])
 
540
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
541
        # The next spills to the 2-len slot
 
542
        builder.add_node(*nodes[4])
 
543
        builder.add_node(*nodes[5])
 
544
        self.assertEqual(0, len(builder._nodes))
 
545
        self.assertIs(None, builder._nodes_by_key)
 
546
        self.assertEqual(2, len(builder._backing_indices))
 
547
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
548
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
549
        # Next spill combines
 
550
        builder.add_node(*nodes[6])
 
551
        builder.add_node(*nodes[7])
 
552
        self.assertEqual(3, len(builder._backing_indices))
 
553
        self.assertEqual(None, builder._backing_indices[0])
 
554
        self.assertEqual(None, builder._backing_indices[1])
 
555
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
556
        # And so forth - counting up in binary.
 
557
        builder.add_node(*nodes[8])
 
558
        builder.add_node(*nodes[9])
 
559
        self.assertEqual(3, len(builder._backing_indices))
 
560
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
561
        self.assertEqual(None, builder._backing_indices[1])
 
562
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
563
        builder.add_node(*nodes[10])
 
564
        builder.add_node(*nodes[11])
 
565
        self.assertEqual(3, len(builder._backing_indices))
 
566
        self.assertEqual(None, builder._backing_indices[0])
 
567
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
568
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
569
        builder.add_node(*nodes[12])
 
570
        # Test that memory and disk are both used for query methods; and that
 
571
        # None is skipped over happily.
 
572
        self.assertEqual([(builder,) + node for node in sorted(nodes[:13])],
 
573
            list(builder.iter_all_entries()))
 
574
        # Two nodes - one memory one disk
 
575
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
 
576
            set(builder.iter_entries([nodes[12][0], nodes[11][0]])))
 
577
        self.assertEqual(13, builder.key_count())
 
578
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
 
579
            set(builder.iter_entries_prefix([nodes[12][0], nodes[11][0]])))
 
580
        builder.add_node(*nodes[13])
 
581
        self.assertEqual(3, len(builder._backing_indices))
 
582
        self.assertEqual(2, builder._backing_indices[0].key_count())
 
583
        self.assertEqual(4, builder._backing_indices[1].key_count())
 
584
        self.assertEqual(8, builder._backing_indices[2].key_count())
 
585
        builder.add_node(*nodes[14])
 
586
        builder.add_node(*nodes[15])
 
587
        self.assertEqual(4, len(builder._backing_indices))
 
588
        self.assertEqual(None, builder._backing_indices[0])
 
589
        self.assertEqual(None, builder._backing_indices[1])
 
590
        self.assertEqual(None, builder._backing_indices[2])
 
591
        self.assertEqual(16, builder._backing_indices[3].key_count())
 
592
        # Now finish, and check we got a correctly ordered tree
 
593
        transport = self.get_transport('')
 
594
        size = transport.put_file('index', builder.finish())
 
595
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
596
        nodes = list(index.iter_all_entries())
 
597
        self.assertEqual(sorted(nodes), nodes)
 
598
        self.assertEqual(16, len(nodes))
 
599
 
 
600
    def test_spill_index_duplicate_key_caught_on_finish(self):
 
601
        builder = btree_index.BTreeBuilder(key_elements=1, spill_at=2)
 
602
        nodes = [node[0:2] for node in self.make_nodes(16, 1, 0)]
 
603
        builder.add_node(*nodes[0])
 
604
        builder.add_node(*nodes[1])
 
605
        builder.add_node(*nodes[0])
 
606
        self.assertRaises(errors.BadIndexDuplicateKey, builder.finish)
 
607
 
 
608
 
 
609
class TestBTreeIndex(BTreeTestCase):
 
610
 
 
611
    def make_index(self, ref_lists=0, key_elements=1, nodes=[]):
 
612
        builder = btree_index.BTreeBuilder(reference_lists=ref_lists,
 
613
            key_elements=key_elements)
 
614
        for key, value, references in nodes:
 
615
            builder.add_node(key, value, references)
 
616
        stream = builder.finish()
 
617
        trans = get_transport('trace+' + self.get_url())
 
618
        size = trans.put_file('index', stream)
 
619
        return btree_index.BTreeGraphIndex(trans, 'index', size)
 
620
 
 
621
    def test_clear_cache(self):
 
622
        nodes = self.make_nodes(160, 2, 2)
 
623
        index = self.make_index(ref_lists=2, key_elements=2, nodes=nodes)
 
624
        self.assertEqual(1, len(list(index.iter_entries([nodes[30][0]]))))
 
625
        self.assertEqual([1, 4], index._row_lengths)
 
626
        self.assertIsNot(None, index._root_node)
 
627
        internal_node_pre_clear = index._internal_node_cache.keys()
 
628
        self.assertTrue(len(index._leaf_node_cache) > 0)
 
629
        index.clear_cache()
 
630
        # We don't touch _root_node or _internal_node_cache, both should be
 
631
        # small, and can save a round trip or two
 
632
        self.assertIsNot(None, index._root_node)
 
633
        # NOTE: We don't want to affect the _internal_node_cache, as we expect
 
634
        #       it will be small, and if we ever do touch this index again, it
 
635
        #       will save round-trips.  This assertion isn't very strong,
 
636
        #       becuase without a 3-level index, we don't have any internal
 
637
        #       nodes cached.
 
638
        self.assertEqual(internal_node_pre_clear,
 
639
                         index._internal_node_cache.keys())
 
640
        self.assertEqual(0, len(index._leaf_node_cache))
 
641
 
 
642
    def test_trivial_constructor(self):
 
643
        transport = get_transport('trace+' + self.get_url(''))
 
644
        index = btree_index.BTreeGraphIndex(transport, 'index', None)
 
645
        # Checks the page size at load, but that isn't logged yet.
 
646
        self.assertEqual([], transport._activity)
 
647
 
 
648
    def test_with_size_constructor(self):
 
649
        transport = get_transport('trace+' + self.get_url(''))
 
650
        index = btree_index.BTreeGraphIndex(transport, 'index', 1)
 
651
        # Checks the page size at load, but that isn't logged yet.
 
652
        self.assertEqual([], transport._activity)
 
653
 
 
654
    def test_empty_key_count_no_size(self):
 
655
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
 
656
        transport = get_transport('trace+' + self.get_url(''))
 
657
        transport.put_file('index', builder.finish())
 
658
        index = btree_index.BTreeGraphIndex(transport, 'index', None)
 
659
        del transport._activity[:]
 
660
        self.assertEqual([], transport._activity)
 
661
        self.assertEqual(0, index.key_count())
 
662
        # The entire index should have been requested (as we generally have the
 
663
        # size available, and doing many small readvs is inappropriate).
 
664
        # We can't tell how much was actually read here, but - check the code.
 
665
        self.assertEqual([('get', 'index')], transport._activity)
 
666
 
 
667
    def test_empty_key_count(self):
 
668
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
 
669
        transport = get_transport('trace+' + self.get_url(''))
 
670
        size = transport.put_file('index', builder.finish())
 
671
        self.assertEqual(72, size)
 
672
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
673
        del transport._activity[:]
 
674
        self.assertEqual([], transport._activity)
 
675
        self.assertEqual(0, index.key_count())
 
676
        # The entire index should have been read, as 4K > size
 
677
        self.assertEqual([('readv', 'index', [(0, 72)], False, None)],
 
678
            transport._activity)
 
679
 
 
680
    def test_non_empty_key_count_2_2(self):
 
681
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
682
        nodes = self.make_nodes(35, 2, 2)
 
683
        for node in nodes:
 
684
            builder.add_node(*node)
 
685
        transport = get_transport('trace+' + self.get_url(''))
 
686
        size = transport.put_file('index', builder.finish())
 
687
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
688
        del transport._activity[:]
 
689
        self.assertEqual([], transport._activity)
 
690
        self.assertEqual(70, index.key_count())
 
691
        # The entire index should have been read, as it is one page long.
 
692
        self.assertEqual([('readv', 'index', [(0, size)], False, None)],
 
693
            transport._activity)
 
694
        self.assertEqual(1173, size)
 
695
 
 
696
    def test__read_nodes_no_size_one_page_reads_once(self):
 
697
        self.make_index(nodes=[(('key',), 'value', ())])
 
698
        trans = get_transport('trace+' + self.get_url())
 
699
        index = btree_index.BTreeGraphIndex(trans, 'index', None)
 
700
        del trans._activity[:]
 
701
        nodes = dict(index._read_nodes([0]))
 
702
        self.assertEqual([0], nodes.keys())
 
703
        node = nodes[0]
 
704
        self.assertEqual([('key',)], node.keys.keys())
 
705
        self.assertEqual([('get', 'index')], trans._activity)
 
706
 
 
707
    def test__read_nodes_no_size_multiple_pages(self):
 
708
        index = self.make_index(2, 2, nodes=self.make_nodes(160, 2, 2))
 
709
        index.key_count()
 
710
        num_pages = index._row_offsets[-1]
 
711
        # Reopen with a traced transport and no size
 
712
        trans = get_transport('trace+' + self.get_url())
 
713
        index = btree_index.BTreeGraphIndex(trans, 'index', None)
 
714
        del trans._activity[:]
 
715
        nodes = dict(index._read_nodes([0]))
 
716
        self.assertEqual(range(num_pages), nodes.keys())
 
717
 
 
718
    def test_2_levels_key_count_2_2(self):
 
719
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
720
        nodes = self.make_nodes(160, 2, 2)
 
721
        for node in nodes:
 
722
            builder.add_node(*node)
 
723
        transport = get_transport('trace+' + self.get_url(''))
 
724
        size = transport.put_file('index', builder.finish())
 
725
        self.assertEqual(17692, size)
 
726
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
727
        del transport._activity[:]
 
728
        self.assertEqual([], transport._activity)
 
729
        self.assertEqual(320, index.key_count())
 
730
        # The entire index should not have been read.
 
731
        self.assertEqual([('readv', 'index', [(0, 4096)], False, None)],
 
732
            transport._activity)
 
733
 
 
734
    def test_validate_one_page(self):
 
735
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
736
        nodes = self.make_nodes(45, 2, 2)
 
737
        for node in nodes:
 
738
            builder.add_node(*node)
 
739
        transport = get_transport('trace+' + self.get_url(''))
 
740
        size = transport.put_file('index', builder.finish())
 
741
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
742
        del transport._activity[:]
 
743
        self.assertEqual([], transport._activity)
 
744
        index.validate()
 
745
        # The entire index should have been read linearly.
 
746
        self.assertEqual([('readv', 'index', [(0, size)], False, None)],
 
747
            transport._activity)
 
748
        self.assertEqual(1488, size)
 
749
 
 
750
    def test_validate_two_pages(self):
 
751
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
752
        nodes = self.make_nodes(80, 2, 2)
 
753
        for node in nodes:
 
754
            builder.add_node(*node)
 
755
        transport = get_transport('trace+' + self.get_url(''))
 
756
        size = transport.put_file('index', builder.finish())
 
757
        # Root page, 2 leaf pages
 
758
        self.assertEqual(9339, size)
 
759
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
760
        del transport._activity[:]
 
761
        self.assertEqual([], transport._activity)
 
762
        index.validate()
 
763
        # The entire index should have been read linearly.
 
764
        self.assertEqual([('readv', 'index', [(0, 4096)], False, None),
 
765
            ('readv', 'index', [(4096, 4096), (8192, 1147)], False, None)],
 
766
            transport._activity)
 
767
        # XXX: TODO: write some badly-ordered nodes, and some pointers-to-wrong
 
768
        # node and make validate find them.
 
769
 
 
770
    def test_eq_ne(self):
 
771
        # two indices are equal when constructed with the same parameters:
 
772
        transport1 = get_transport('trace+' + self.get_url(''))
 
773
        transport2 = get_transport(self.get_url(''))
 
774
        self.assertTrue(
 
775
            btree_index.BTreeGraphIndex(transport1, 'index', None) ==
 
776
            btree_index.BTreeGraphIndex(transport1, 'index', None))
 
777
        self.assertTrue(
 
778
            btree_index.BTreeGraphIndex(transport1, 'index', 20) ==
 
779
            btree_index.BTreeGraphIndex(transport1, 'index', 20))
 
780
        self.assertFalse(
 
781
            btree_index.BTreeGraphIndex(transport1, 'index', 20) ==
 
782
            btree_index.BTreeGraphIndex(transport2, 'index', 20))
 
783
        self.assertFalse(
 
784
            btree_index.BTreeGraphIndex(transport1, 'inde1', 20) ==
 
785
            btree_index.BTreeGraphIndex(transport1, 'inde2', 20))
 
786
        self.assertFalse(
 
787
            btree_index.BTreeGraphIndex(transport1, 'index', 10) ==
 
788
            btree_index.BTreeGraphIndex(transport1, 'index', 20))
 
789
        self.assertFalse(
 
790
            btree_index.BTreeGraphIndex(transport1, 'index', None) !=
 
791
            btree_index.BTreeGraphIndex(transport1, 'index', None))
 
792
        self.assertFalse(
 
793
            btree_index.BTreeGraphIndex(transport1, 'index', 20) !=
 
794
            btree_index.BTreeGraphIndex(transport1, 'index', 20))
 
795
        self.assertTrue(
 
796
            btree_index.BTreeGraphIndex(transport1, 'index', 20) !=
 
797
            btree_index.BTreeGraphIndex(transport2, 'index', 20))
 
798
        self.assertTrue(
 
799
            btree_index.BTreeGraphIndex(transport1, 'inde1', 20) !=
 
800
            btree_index.BTreeGraphIndex(transport1, 'inde2', 20))
 
801
        self.assertTrue(
 
802
            btree_index.BTreeGraphIndex(transport1, 'index', 10) !=
 
803
            btree_index.BTreeGraphIndex(transport1, 'index', 20))
 
804
 
 
805
    def test_iter_all_only_root_no_size(self):
 
806
        self.make_index(nodes=[(('key',), 'value', ())])
 
807
        trans = get_transport('trace+' + self.get_url(''))
 
808
        index = btree_index.BTreeGraphIndex(trans, 'index', None)
 
809
        del trans._activity[:]
 
810
        self.assertEqual([(('key',), 'value')],
 
811
                         [x[1:] for x in index.iter_all_entries()])
 
812
        self.assertEqual([('get', 'index')], trans._activity)
 
813
 
 
814
    def test_iter_all_entries_reads(self):
 
815
        # iterating all entries reads the header, then does a linear
 
816
        # read.
 
817
        self.shrink_page_size()
 
818
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
819
        # 20k nodes is enough to create a two internal nodes on the second
 
820
        # level, with a 2K page size
 
821
        nodes = self.make_nodes(10000, 2, 2)
 
822
        for node in nodes:
 
823
            builder.add_node(*node)
 
824
        transport = get_transport('trace+' + self.get_url(''))
 
825
        size = transport.put_file('index', builder.finish())
 
826
        self.assertEqual(1303220, size, 'number of expected bytes in the'
 
827
                                        ' output changed')
 
828
        page_size = btree_index._PAGE_SIZE
 
829
        del builder
 
830
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
831
        del transport._activity[:]
 
832
        self.assertEqual([], transport._activity)
 
833
        found_nodes = self.time(list, index.iter_all_entries())
 
834
        bare_nodes = []
 
835
        for node in found_nodes:
 
836
            self.assertTrue(node[0] is index)
 
837
            bare_nodes.append(node[1:])
 
838
        self.assertEqual(3, len(index._row_lengths),
 
839
            "Not enough rows: %r" % index._row_lengths)
 
840
        # Should be as long as the nodes we supplied
 
841
        self.assertEqual(20000, len(found_nodes))
 
842
        # Should have the same content
 
843
        self.assertEqual(set(nodes), set(bare_nodes))
 
844
        # Should have done linear scan IO up the index, ignoring
 
845
        # the internal nodes:
 
846
        # The entire index should have been read
 
847
        total_pages = sum(index._row_lengths)
 
848
        self.assertEqual(total_pages, index._row_offsets[-1])
 
849
        self.assertEqual(1303220, size)
 
850
        # The start of the leaves
 
851
        first_byte = index._row_offsets[-2] * page_size
 
852
        readv_request = []
 
853
        for offset in range(first_byte, size, page_size):
 
854
            readv_request.append((offset, page_size))
 
855
        # The last page is truncated
 
856
        readv_request[-1] = (readv_request[-1][0], 1303220 % page_size)
 
857
        expected = [('readv', 'index', [(0, page_size)], False, None),
 
858
             ('readv',  'index', readv_request, False, None)]
 
859
        if expected != transport._activity:
 
860
            self.assertEqualDiff(pprint.pformat(expected),
 
861
                                 pprint.pformat(transport._activity))
 
862
 
 
863
    def _test_iter_entries_references_resolved(self):
 
864
        index = self.make_index(1, nodes=[
 
865
            (('name', ), 'data', ([('ref', ), ('ref', )], )),
 
866
            (('ref', ), 'refdata', ([], ))])
 
867
        self.assertEqual(set([(index, ('name', ), 'data', ((('ref',),('ref',)),)),
 
868
            (index, ('ref', ), 'refdata', ((), ))]),
 
869
            set(index.iter_entries([('name',), ('ref',)])))
 
870
 
 
871
    def test_iter_entries_references_2_refs_resolved(self):
 
872
        # iterating some entries reads just the pages needed. For now, to
 
873
        # get it working and start measuring, only 4K pages are read.
 
874
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
 
875
        # 80 nodes is enough to create a two-level index.
 
876
        nodes = self.make_nodes(160, 2, 2)
 
877
        for node in nodes:
 
878
            builder.add_node(*node)
 
879
        transport = get_transport('trace+' + self.get_url(''))
 
880
        size = transport.put_file('index', builder.finish())
 
881
        del builder
 
882
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
 
883
        del transport._activity[:]
 
884
        self.assertEqual([], transport._activity)
 
885
        # search for one key
 
886
        found_nodes = list(index.iter_entries([nodes[30][0]]))
 
887
        bare_nodes = []
 
888
        for node in found_nodes:
 
889
            self.assertTrue(node[0] is index)
 
890
            bare_nodes.append(node[1:])
 
891
        # Should be as long as the nodes we supplied
 
892
        self.assertEqual(1, len(found_nodes))
 
893
        # Should have the same content
 
894
        self.assertEqual(nodes[30], bare_nodes[0])
 
895
        # Should have read the root node, then one leaf page:
 
896
        self.assertEqual([('readv', 'index', [(0, 4096)], False, None),
 
897
             ('readv',  'index', [(8192, 4096), ], False, None)],
 
898
            transport._activity)
 
899
 
 
900
    def test_iter_key_prefix_1_element_key_None(self):
 
901
        index = self.make_index()
 
902
        self.assertRaises(errors.BadIndexKey, list,
 
903
            index.iter_entries_prefix([(None, )]))
 
904
 
 
905
    def test_iter_key_prefix_wrong_length(self):
 
906
        index = self.make_index()
 
907
        self.assertRaises(errors.BadIndexKey, list,
 
908
            index.iter_entries_prefix([('foo', None)]))
 
909
        index = self.make_index(key_elements=2)
 
910
        self.assertRaises(errors.BadIndexKey, list,
 
911
            index.iter_entries_prefix([('foo', )]))
 
912
        self.assertRaises(errors.BadIndexKey, list,
 
913
            index.iter_entries_prefix([('foo', None, None)]))
 
914
 
 
915
    def test_iter_key_prefix_1_key_element_no_refs(self):
 
916
        index = self.make_index( nodes=[
 
917
            (('name', ), 'data', ()),
 
918
            (('ref', ), 'refdata', ())])
 
919
        self.assertEqual(set([(index, ('name', ), 'data'),
 
920
            (index, ('ref', ), 'refdata')]),
 
921
            set(index.iter_entries_prefix([('name', ), ('ref', )])))
 
922
 
 
923
    def test_iter_key_prefix_1_key_element_refs(self):
 
924
        index = self.make_index(1, nodes=[
 
925
            (('name', ), 'data', ([('ref', )], )),
 
926
            (('ref', ), 'refdata', ([], ))])
 
927
        self.assertEqual(set([(index, ('name', ), 'data', ((('ref',),),)),
 
928
            (index, ('ref', ), 'refdata', ((), ))]),
 
929
            set(index.iter_entries_prefix([('name', ), ('ref', )])))
 
930
 
 
931
    def test_iter_key_prefix_2_key_element_no_refs(self):
 
932
        index = self.make_index(key_elements=2, nodes=[
 
933
            (('name', 'fin1'), 'data', ()),
 
934
            (('name', 'fin2'), 'beta', ()),
 
935
            (('ref', 'erence'), 'refdata', ())])
 
936
        self.assertEqual(set([(index, ('name', 'fin1'), 'data'),
 
937
            (index, ('ref', 'erence'), 'refdata')]),
 
938
            set(index.iter_entries_prefix([('name', 'fin1'), ('ref', 'erence')])))
 
939
        self.assertEqual(set([(index, ('name', 'fin1'), 'data'),
 
940
            (index, ('name', 'fin2'), 'beta')]),
 
941
            set(index.iter_entries_prefix([('name', None)])))
 
942
 
 
943
    def test_iter_key_prefix_2_key_element_refs(self):
 
944
        index = self.make_index(1, key_elements=2, nodes=[
 
945
            (('name', 'fin1'), 'data', ([('ref', 'erence')], )),
 
946
            (('name', 'fin2'), 'beta', ([], )),
 
947
            (('ref', 'erence'), 'refdata', ([], ))])
 
948
        self.assertEqual(set([(index, ('name', 'fin1'), 'data', ((('ref', 'erence'),),)),
 
949
            (index, ('ref', 'erence'), 'refdata', ((), ))]),
 
950
            set(index.iter_entries_prefix([('name', 'fin1'), ('ref', 'erence')])))
 
951
        self.assertEqual(set([(index, ('name', 'fin1'), 'data', ((('ref', 'erence'),),)),
 
952
            (index, ('name', 'fin2'), 'beta', ((), ))]),
 
953
            set(index.iter_entries_prefix([('name', None)])))
 
954
 
 
955
    # XXX: external_references tests are duplicated in test_index.  We
 
956
    # probably should have per_graph_index tests...
 
957
    def test_external_references_no_refs(self):
 
958
        index = self.make_index(ref_lists=0, nodes=[])
 
959
        self.assertRaises(ValueError, index.external_references, 0)
 
960
 
 
961
    def test_external_references_no_results(self):
 
962
        index = self.make_index(ref_lists=1, nodes=[
 
963
            (('key',), 'value', ([],))])
 
964
        self.assertEqual(set(), index.external_references(0))
 
965
 
 
966
    def test_external_references_missing_ref(self):
 
967
        missing_key = ('missing',)
 
968
        index = self.make_index(ref_lists=1, nodes=[
 
969
            (('key',), 'value', ([missing_key],))])
 
970
        self.assertEqual(set([missing_key]), index.external_references(0))
 
971
 
 
972
    def test_external_references_multiple_ref_lists(self):
 
973
        missing_key = ('missing',)
 
974
        index = self.make_index(ref_lists=2, nodes=[
 
975
            (('key',), 'value', ([], [missing_key]))])
 
976
        self.assertEqual(set([]), index.external_references(0))
 
977
        self.assertEqual(set([missing_key]), index.external_references(1))
 
978
 
 
979
    def test_external_references_two_records(self):
 
980
        index = self.make_index(ref_lists=1, nodes=[
 
981
            (('key-1',), 'value', ([('key-2',)],)),
 
982
            (('key-2',), 'value', ([],)),
 
983
            ])
 
984
        self.assertEqual(set([]), index.external_references(0))
 
985
 
 
986
    def test__find_ancestors_one_page(self):
 
987
        key1 = ('key-1',)
 
988
        key2 = ('key-2',)
 
989
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
 
990
            (key1, 'value', ([key2],)),
 
991
            (key2, 'value', ([],)),
 
992
            ])
 
993
        parent_map = {}
 
994
        missing_keys = set()
 
995
        search_keys = index._find_ancestors([key1], 0, parent_map, missing_keys)
 
996
        self.assertEqual({key1: (key2,), key2: ()}, parent_map)
 
997
        self.assertEqual(set(), missing_keys)
 
998
        self.assertEqual(set(), search_keys)
 
999
 
 
1000
    def test__find_ancestors_one_page_w_missing(self):
 
1001
        key1 = ('key-1',)
 
1002
        key2 = ('key-2',)
 
1003
        key3 = ('key-3',)
 
1004
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
 
1005
            (key1, 'value', ([key2],)),
 
1006
            (key2, 'value', ([],)),
 
1007
            ])
 
1008
        parent_map = {}
 
1009
        missing_keys = set()
 
1010
        search_keys = index._find_ancestors([key2, key3], 0, parent_map,
 
1011
                                            missing_keys)
 
1012
        self.assertEqual({key2: ()}, parent_map)
 
1013
        # we know that key3 is missing because we read the page that it would
 
1014
        # otherwise be on
 
1015
        self.assertEqual(set([key3]), missing_keys)
 
1016
        self.assertEqual(set(), search_keys)
 
1017
 
 
1018
    def test__find_ancestors_one_parent_missing(self):
 
1019
        key1 = ('key-1',)
 
1020
        key2 = ('key-2',)
 
1021
        key3 = ('key-3',)
 
1022
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
 
1023
            (key1, 'value', ([key2],)),
 
1024
            (key2, 'value', ([key3],)),
 
1025
            ])
 
1026
        parent_map = {}
 
1027
        missing_keys = set()
 
1028
        search_keys = index._find_ancestors([key1], 0, parent_map,
 
1029
                                            missing_keys)
 
1030
        self.assertEqual({key1: (key2,), key2: (key3,)}, parent_map)
 
1031
        self.assertEqual(set(), missing_keys)
 
1032
        # all we know is that key3 wasn't present on the page we were reading
 
1033
        # but if you look, the last key is key2 which comes before key3, so we
 
1034
        # don't know whether key3 would land on this page or not.
 
1035
        self.assertEqual(set([key3]), search_keys)
 
1036
        search_keys = index._find_ancestors(search_keys, 0, parent_map,
 
1037
                                            missing_keys)
 
1038
        # passing it back in, we are sure it is 'missing'
 
1039
        self.assertEqual({key1: (key2,), key2: (key3,)}, parent_map)
 
1040
        self.assertEqual(set([key3]), missing_keys)
 
1041
        self.assertEqual(set([]), search_keys)
 
1042
 
 
1043
    def test__find_ancestors_dont_search_known(self):
 
1044
        key1 = ('key-1',)
 
1045
        key2 = ('key-2',)
 
1046
        key3 = ('key-3',)
 
1047
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
 
1048
            (key1, 'value', ([key2],)),
 
1049
            (key2, 'value', ([key3],)),
 
1050
            (key3, 'value', ([],)),
 
1051
            ])
 
1052
        # We already know about key2, so we won't try to search for key3
 
1053
        parent_map = {key2: (key3,)}
 
1054
        missing_keys = set()
 
1055
        search_keys = index._find_ancestors([key1], 0, parent_map,
 
1056
                                            missing_keys)
 
1057
        self.assertEqual({key1: (key2,), key2: (key3,)}, parent_map)
 
1058
        self.assertEqual(set(), missing_keys)
 
1059
        self.assertEqual(set(), search_keys)
 
1060
 
 
1061
    def test__find_ancestors_multiple_pages(self):
 
1062
        # We need to use enough keys that we actually cause a split
 
1063
        start_time = 1249671539
 
1064
        email = "joebob@example.com"
 
1065
        nodes = []
 
1066
        ref_lists = ((),)
 
1067
        rev_keys = []
 
1068
        for i in xrange(400):
 
1069
            rev_id = '%s-%s-%s' % (email,
 
1070
                                   osutils.compact_date(start_time + i),
 
1071
                                   osutils.rand_chars(16))
 
1072
            rev_key = (rev_id,)
 
1073
            nodes.append((rev_key, 'value', ref_lists))
 
1074
            # We have a ref 'list' of length 1, with a list of parents, with 1
 
1075
            # parent which is a key
 
1076
            ref_lists = ((rev_key,),)
 
1077
            rev_keys.append(rev_key)
 
1078
        index = self.make_index(ref_lists=1, key_elements=1, nodes=nodes)
 
1079
        self.assertEqual(400, index.key_count())
 
1080
        self.assertEqual(3, len(index._row_offsets))
 
1081
        nodes = dict(index._read_nodes([1, 2]))
 
1082
        l1 = nodes[1]
 
1083
        l2 = nodes[2]
 
1084
        min_l2_key = l2.min_key
 
1085
        max_l1_key = l1.max_key
 
1086
        self.assertTrue(max_l1_key < min_l2_key)
 
1087
        parents_min_l2_key = l2.keys[min_l2_key][1][0]
 
1088
        self.assertEqual((l1.max_key,), parents_min_l2_key)
 
1089
        # Now, whatever key we select that would fall on the second page,
 
1090
        # should give us all the parents until the page break
 
1091
        key_idx = rev_keys.index(min_l2_key)
 
1092
        next_key = rev_keys[key_idx+1]
 
1093
        # So now when we get the parent map, we should get the key we are
 
1094
        # looking for, min_l2_key, and then a reference to go look for the
 
1095
        # parent of that key
 
1096
        parent_map = {}
 
1097
        missing_keys = set()
 
1098
        search_keys = index._find_ancestors([next_key], 0, parent_map,
 
1099
                                            missing_keys)
 
1100
        self.assertEqual([min_l2_key, next_key], sorted(parent_map))
 
1101
        self.assertEqual(set(), missing_keys)
 
1102
        self.assertEqual(set([max_l1_key]), search_keys)
 
1103
        parent_map = {}
 
1104
        search_keys = index._find_ancestors([max_l1_key], 0, parent_map,
 
1105
                                            missing_keys)
 
1106
        self.assertEqual(sorted(l1.keys), sorted(parent_map))
 
1107
        self.assertEqual(set(), missing_keys)
 
1108
        self.assertEqual(set(), search_keys)
 
1109
 
 
1110
    def test__find_ancestors_empty_index(self):
 
1111
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[])
 
1112
        parent_map = {}
 
1113
        missing_keys = set()
 
1114
        search_keys = index._find_ancestors([('one',), ('two',)], 0, parent_map,
 
1115
                                            missing_keys)
 
1116
        self.assertEqual(set(), search_keys)
 
1117
        self.assertEqual({}, parent_map)
 
1118
        self.assertEqual(set([('one',), ('two',)]), missing_keys)
 
1119
 
 
1120
    def test_supports_unlimited_cache(self):
 
1121
        builder = btree_index.BTreeBuilder(reference_lists=0, key_elements=1)
 
1122
        # We need enough nodes to cause a page split (so we have both an
 
1123
        # internal node and a couple leaf nodes. 500 seems to be enough.)
 
1124
        nodes = self.make_nodes(500, 1, 0)
 
1125
        for node in nodes:
 
1126
            builder.add_node(*node)
 
1127
        stream = builder.finish()
 
1128
        trans = get_transport(self.get_url())
 
1129
        size = trans.put_file('index', stream)
 
1130
        index = btree_index.BTreeGraphIndex(trans, 'index', size)
 
1131
        self.assertEqual(500, index.key_count())
 
1132
        # We have an internal node
 
1133
        self.assertEqual(2, len(index._row_lengths))
 
1134
        # We have at least 2 leaf nodes
 
1135
        self.assertTrue(index._row_lengths[-1] >= 2)
 
1136
        self.assertIsInstance(index._leaf_node_cache, lru_cache.LRUCache)
 
1137
        self.assertEqual(btree_index._NODE_CACHE_SIZE,
 
1138
                         index._leaf_node_cache._max_cache)
 
1139
        self.assertIsInstance(index._internal_node_cache, fifo_cache.FIFOCache)
 
1140
        self.assertEqual(100, index._internal_node_cache._max_cache)
 
1141
        # No change if unlimited_cache=False is passed
 
1142
        index = btree_index.BTreeGraphIndex(trans, 'index', size,
 
1143
                                            unlimited_cache=False)
 
1144
        self.assertIsInstance(index._leaf_node_cache, lru_cache.LRUCache)
 
1145
        self.assertEqual(btree_index._NODE_CACHE_SIZE,
 
1146
                         index._leaf_node_cache._max_cache)
 
1147
        self.assertIsInstance(index._internal_node_cache, fifo_cache.FIFOCache)
 
1148
        self.assertEqual(100, index._internal_node_cache._max_cache)
 
1149
        index = btree_index.BTreeGraphIndex(trans, 'index', size,
 
1150
                                            unlimited_cache=True)
 
1151
        self.assertIsInstance(index._leaf_node_cache, dict)
 
1152
        self.assertIs(type(index._internal_node_cache), dict)
 
1153
        # Exercise the lookup code
 
1154
        entries = set(index.iter_entries([n[0] for n in nodes]))
 
1155
        self.assertEqual(500, len(entries))
 
1156
 
 
1157
 
 
1158
class TestBTreeNodes(BTreeTestCase):
 
1159
 
 
1160
    def restore_parser(self):
 
1161
        btree_index._btree_serializer = self.saved_parser
 
1162
 
 
1163
    def setUp(self):
 
1164
        BTreeTestCase.setUp(self)
 
1165
        self.saved_parser = btree_index._btree_serializer
 
1166
        self.addCleanup(self.restore_parser)
 
1167
        btree_index._btree_serializer = self.parse_btree
 
1168
 
 
1169
    def test_LeafNode_1_0(self):
 
1170
        node_bytes = ("type=leaf\n"
 
1171
            "0000000000000000000000000000000000000000\x00\x00value:0\n"
 
1172
            "1111111111111111111111111111111111111111\x00\x00value:1\n"
 
1173
            "2222222222222222222222222222222222222222\x00\x00value:2\n"
 
1174
            "3333333333333333333333333333333333333333\x00\x00value:3\n"
 
1175
            "4444444444444444444444444444444444444444\x00\x00value:4\n")
 
1176
        node = btree_index._LeafNode(node_bytes, 1, 0)
 
1177
        # We do direct access, or don't care about order, to leaf nodes most of
 
1178
        # the time, so a dict is useful:
 
1179
        self.assertEqual({
 
1180
            ("0000000000000000000000000000000000000000",): ("value:0", ()),
 
1181
            ("1111111111111111111111111111111111111111",): ("value:1", ()),
 
1182
            ("2222222222222222222222222222222222222222",): ("value:2", ()),
 
1183
            ("3333333333333333333333333333333333333333",): ("value:3", ()),
 
1184
            ("4444444444444444444444444444444444444444",): ("value:4", ()),
 
1185
            }, node.keys)
 
1186
 
 
1187
    def test_LeafNode_2_2(self):
 
1188
        node_bytes = ("type=leaf\n"
 
1189
            "00\x0000\x00\t00\x00ref00\x00value:0\n"
 
1190
            "00\x0011\x0000\x00ref00\t00\x00ref00\r01\x00ref01\x00value:1\n"
 
1191
            "11\x0033\x0011\x00ref22\t11\x00ref22\r11\x00ref22\x00value:3\n"
 
1192
            "11\x0044\x00\t11\x00ref00\x00value:4\n"
 
1193
            ""
 
1194
            )
 
1195
        node = btree_index._LeafNode(node_bytes, 2, 2)
 
1196
        # We do direct access, or don't care about order, to leaf nodes most of
 
1197
        # the time, so a dict is useful:
 
1198
        self.assertEqual({
 
1199
            ('00', '00'): ('value:0', ((), (('00', 'ref00'),))),
 
1200
            ('00', '11'): ('value:1',
 
1201
                ((('00', 'ref00'),), (('00', 'ref00'), ('01', 'ref01')))),
 
1202
            ('11', '33'): ('value:3',
 
1203
                ((('11', 'ref22'),), (('11', 'ref22'), ('11', 'ref22')))),
 
1204
            ('11', '44'): ('value:4', ((), (('11', 'ref00'),)))
 
1205
            }, node.keys)
 
1206
 
 
1207
    def test_InternalNode_1(self):
 
1208
        node_bytes = ("type=internal\n"
 
1209
            "offset=1\n"
 
1210
            "0000000000000000000000000000000000000000\n"
 
1211
            "1111111111111111111111111111111111111111\n"
 
1212
            "2222222222222222222222222222222222222222\n"
 
1213
            "3333333333333333333333333333333333333333\n"
 
1214
            "4444444444444444444444444444444444444444\n"
 
1215
            )
 
1216
        node = btree_index._InternalNode(node_bytes)
 
1217
        # We want to bisect to find the right children from this node, so a
 
1218
        # vector is most useful.
 
1219
        self.assertEqual([
 
1220
            ("0000000000000000000000000000000000000000",),
 
1221
            ("1111111111111111111111111111111111111111",),
 
1222
            ("2222222222222222222222222222222222222222",),
 
1223
            ("3333333333333333333333333333333333333333",),
 
1224
            ("4444444444444444444444444444444444444444",),
 
1225
            ], node.keys)
 
1226
        self.assertEqual(1, node.offset)
 
1227
 
 
1228
    def test_LeafNode_2_2(self):
 
1229
        node_bytes = ("type=leaf\n"
 
1230
            "00\x0000\x00\t00\x00ref00\x00value:0\n"
 
1231
            "00\x0011\x0000\x00ref00\t00\x00ref00\r01\x00ref01\x00value:1\n"
 
1232
            "11\x0033\x0011\x00ref22\t11\x00ref22\r11\x00ref22\x00value:3\n"
 
1233
            "11\x0044\x00\t11\x00ref00\x00value:4\n"
 
1234
            ""
 
1235
            )
 
1236
        node = btree_index._LeafNode(node_bytes, 2, 2)
 
1237
        # We do direct access, or don't care about order, to leaf nodes most of
 
1238
        # the time, so a dict is useful:
 
1239
        self.assertEqual({
 
1240
            ('00', '00'): ('value:0', ((), (('00', 'ref00'),))),
 
1241
            ('00', '11'): ('value:1',
 
1242
                ((('00', 'ref00'),), (('00', 'ref00'), ('01', 'ref01')))),
 
1243
            ('11', '33'): ('value:3',
 
1244
                ((('11', 'ref22'),), (('11', 'ref22'), ('11', 'ref22')))),
 
1245
            ('11', '44'): ('value:4', ((), (('11', 'ref00'),)))
 
1246
            }, node.keys)
 
1247
 
 
1248
    def assertFlattened(self, expected, key, value, refs):
 
1249
        flat_key, flat_line = self.parse_btree._flatten_node(
 
1250
            (None, key, value, refs), bool(refs))
 
1251
        self.assertEqual('\x00'.join(key), flat_key)
 
1252
        self.assertEqual(expected, flat_line)
 
1253
 
 
1254
    def test__flatten_node(self):
 
1255
        self.assertFlattened('key\0\0value\n', ('key',), 'value', [])
 
1256
        self.assertFlattened('key\0tuple\0\0value str\n',
 
1257
                             ('key', 'tuple'), 'value str', [])
 
1258
        self.assertFlattened('key\0tuple\0triple\0\0value str\n',
 
1259
                             ('key', 'tuple', 'triple'), 'value str', [])
 
1260
        self.assertFlattened('k\0t\0s\0ref\0value str\n',
 
1261
                             ('k', 't', 's'), 'value str', [[('ref',)]])
 
1262
        self.assertFlattened('key\0tuple\0ref\0key\0value str\n',
 
1263
                             ('key', 'tuple'), 'value str', [[('ref', 'key')]])
 
1264
        self.assertFlattened("00\x0000\x00\t00\x00ref00\x00value:0\n",
 
1265
            ('00', '00'), 'value:0', ((), (('00', 'ref00'),)))
 
1266
        self.assertFlattened(
 
1267
            "00\x0011\x0000\x00ref00\t00\x00ref00\r01\x00ref01\x00value:1\n",
 
1268
            ('00', '11'), 'value:1',
 
1269
                ((('00', 'ref00'),), (('00', 'ref00'), ('01', 'ref01'))))
 
1270
        self.assertFlattened(
 
1271
            "11\x0033\x0011\x00ref22\t11\x00ref22\r11\x00ref22\x00value:3\n",
 
1272
            ('11', '33'), 'value:3',
 
1273
                ((('11', 'ref22'),), (('11', 'ref22'), ('11', 'ref22'))))
 
1274
        self.assertFlattened(
 
1275
            "11\x0044\x00\t11\x00ref00\x00value:4\n",
 
1276
            ('11', '44'), 'value:4', ((), (('11', 'ref00'),)))
 
1277
 
 
1278
 
 
1279
class TestCompiledBtree(tests.TestCase):
 
1280
 
 
1281
    def test_exists(self):
 
1282
        # This is just to let the user know if they don't have the feature
 
1283
        # available
 
1284
        self.requireFeature(compiled_btreeparser_feature)
 
1285
 
 
1286
 
 
1287
class TestMultiBisectRight(tests.TestCase):
 
1288
 
 
1289
    def assertMultiBisectRight(self, offsets, search_keys, fixed_keys):
 
1290
        self.assertEqual(offsets,
 
1291
                         btree_index.BTreeGraphIndex._multi_bisect_right(
 
1292
                            search_keys, fixed_keys))
 
1293
 
 
1294
    def test_after(self):
 
1295
        self.assertMultiBisectRight([(1, ['b'])], ['b'], ['a'])
 
1296
        self.assertMultiBisectRight([(3, ['e', 'f', 'g'])],
 
1297
                                    ['e', 'f', 'g'], ['a', 'b', 'c'])
 
1298
 
 
1299
    def test_before(self):
 
1300
        self.assertMultiBisectRight([(0, ['a'])], ['a'], ['b'])
 
1301
        self.assertMultiBisectRight([(0, ['a', 'b', 'c', 'd'])],
 
1302
                                    ['a', 'b', 'c', 'd'], ['e', 'f', 'g'])
 
1303
 
 
1304
    def test_exact(self):
 
1305
        self.assertMultiBisectRight([(1, ['a'])], ['a'], ['a'])
 
1306
        self.assertMultiBisectRight([(1, ['a']), (2, ['b'])], ['a', 'b'], ['a', 'b'])
 
1307
        self.assertMultiBisectRight([(1, ['a']), (3, ['c'])],
 
1308
                                    ['a', 'c'], ['a', 'b', 'c'])
 
1309
 
 
1310
    def test_inbetween(self):
 
1311
        self.assertMultiBisectRight([(1, ['b'])], ['b'], ['a', 'c'])
 
1312
        self.assertMultiBisectRight([(1, ['b', 'c', 'd']), (2, ['f', 'g'])],
 
1313
                                    ['b', 'c', 'd', 'f', 'g'], ['a', 'e', 'h'])
 
1314
 
 
1315
    def test_mixed(self):
 
1316
        self.assertMultiBisectRight([(0, ['a', 'b']), (2, ['d', 'e']),
 
1317
                                     (4, ['g', 'h'])],
 
1318
                                    ['a', 'b', 'd', 'e', 'g', 'h'],
 
1319
                                    ['c', 'd', 'f', 'g'])
 
1320
 
 
1321
 
 
1322
class TestExpandOffsets(tests.TestCase):
 
1323
 
 
1324
    def make_index(self, size, recommended_pages=None):
 
1325
        """Make an index with a generic size.
 
1326
 
 
1327
        This doesn't actually create anything on disk, it just primes a
 
1328
        BTreeGraphIndex with the recommended information.
 
1329
        """
 
1330
        index = btree_index.BTreeGraphIndex(get_transport('memory:///'),
 
1331
                                            'test-index', size=size)
 
1332
        if recommended_pages is not None:
 
1333
            index._recommended_pages = recommended_pages
 
1334
        return index
 
1335
 
 
1336
    def set_cached_offsets(self, index, cached_offsets):
 
1337
        """Monkeypatch to give a canned answer for _get_offsets_for...()."""
 
1338
        def _get_offsets_to_cached_pages():
 
1339
            cached = set(cached_offsets)
 
1340
            return cached
 
1341
        index._get_offsets_to_cached_pages = _get_offsets_to_cached_pages
 
1342
 
 
1343
    def prepare_index(self, index, node_ref_lists, key_length, key_count,
 
1344
                      row_lengths, cached_offsets):
 
1345
        """Setup the BTreeGraphIndex with some pre-canned information."""
 
1346
        index.node_ref_lists = node_ref_lists
 
1347
        index._key_length = key_length
 
1348
        index._key_count = key_count
 
1349
        index._row_lengths = row_lengths
 
1350
        index._compute_row_offsets()
 
1351
        index._root_node = btree_index._InternalNode('internal\noffset=0\n')
 
1352
        self.set_cached_offsets(index, cached_offsets)
 
1353
 
 
1354
    def make_100_node_index(self):
 
1355
        index = self.make_index(4096*100, 6)
 
1356
        # Consider we've already made a single request at the middle
 
1357
        self.prepare_index(index, node_ref_lists=0, key_length=1,
 
1358
                           key_count=1000, row_lengths=[1, 99],
 
1359
                           cached_offsets=[0, 50])
 
1360
        return index
 
1361
 
 
1362
    def make_1000_node_index(self):
 
1363
        index = self.make_index(4096*1000, 6)
 
1364
        # Pretend we've already made a single request in the middle
 
1365
        self.prepare_index(index, node_ref_lists=0, key_length=1,
 
1366
                           key_count=90000, row_lengths=[1, 9, 990],
 
1367
                           cached_offsets=[0, 5, 500])
 
1368
        return index
 
1369
 
 
1370
    def assertNumPages(self, expected_pages, index, size):
 
1371
        index._size = size
 
1372
        self.assertEqual(expected_pages, index._compute_total_pages_in_index())
 
1373
 
 
1374
    def assertExpandOffsets(self, expected, index, offsets):
 
1375
        self.assertEqual(expected, index._expand_offsets(offsets),
 
1376
                         'We did not get the expected value after expanding'
 
1377
                         ' %s' % (offsets,))
 
1378
 
 
1379
    def test_default_recommended_pages(self):
 
1380
        index = self.make_index(None)
 
1381
        # local transport recommends 4096 byte reads, which is 1 page
 
1382
        self.assertEqual(1, index._recommended_pages)
 
1383
 
 
1384
    def test__compute_total_pages_in_index(self):
 
1385
        index = self.make_index(None)
 
1386
        self.assertNumPages(1, index, 1024)
 
1387
        self.assertNumPages(1, index, 4095)
 
1388
        self.assertNumPages(1, index, 4096)
 
1389
        self.assertNumPages(2, index, 4097)
 
1390
        self.assertNumPages(2, index, 8192)
 
1391
        self.assertNumPages(76, index, 4096*75 + 10)
 
1392
 
 
1393
    def test__find_layer_start_and_stop(self):
 
1394
        index = self.make_1000_node_index()
 
1395
        self.assertEqual((0, 1), index._find_layer_first_and_end(0))
 
1396
        self.assertEqual((1, 10), index._find_layer_first_and_end(1))
 
1397
        self.assertEqual((1, 10), index._find_layer_first_and_end(9))
 
1398
        self.assertEqual((10, 1000), index._find_layer_first_and_end(10))
 
1399
        self.assertEqual((10, 1000), index._find_layer_first_and_end(99))
 
1400
        self.assertEqual((10, 1000), index._find_layer_first_and_end(999))
 
1401
 
 
1402
    def test_unknown_size(self):
 
1403
        # We should not expand if we don't know the file size
 
1404
        index = self.make_index(None, 10)
 
1405
        self.assertExpandOffsets([0], index, [0])
 
1406
        self.assertExpandOffsets([1, 4, 9], index, [1, 4, 9])
 
1407
 
 
1408
    def test_more_than_recommended(self):
 
1409
        index = self.make_index(4096*100, 2)
 
1410
        self.assertExpandOffsets([1, 10], index, [1, 10])
 
1411
        self.assertExpandOffsets([1, 10, 20], index, [1, 10, 20])
 
1412
 
 
1413
    def test_read_all_from_root(self):
 
1414
        index = self.make_index(4096*10, 20)
 
1415
        self.assertExpandOffsets(range(10), index, [0])
 
1416
 
 
1417
    def test_read_all_when_cached(self):
 
1418
        # We've read enough that we can grab all the rest in a single request
 
1419
        index = self.make_index(4096*10, 5)
 
1420
        self.prepare_index(index, node_ref_lists=0, key_length=1,
 
1421
                           key_count=1000, row_lengths=[1, 9],
 
1422
                           cached_offsets=[0, 1, 2, 5, 6])
 
1423
        # It should fill the remaining nodes, regardless of the one requested
 
1424
        self.assertExpandOffsets([3, 4, 7, 8, 9], index, [3])
 
1425
        self.assertExpandOffsets([3, 4, 7, 8, 9], index, [8])
 
1426
        self.assertExpandOffsets([3, 4, 7, 8, 9], index, [9])
 
1427
 
 
1428
    def test_no_root_node(self):
 
1429
        index = self.make_index(4096*10, 5)
 
1430
        self.assertExpandOffsets([0], index, [0])
 
1431
 
 
1432
    def test_include_neighbors(self):
 
1433
        index = self.make_100_node_index()
 
1434
        # We expand in both directions, until we have at least 'recommended'
 
1435
        # pages
 
1436
        self.assertExpandOffsets([9, 10, 11, 12, 13, 14, 15], index, [12])
 
1437
        self.assertExpandOffsets([88, 89, 90, 91, 92, 93, 94], index, [91])
 
1438
        # If we hit an 'edge' we continue in the other direction
 
1439
        self.assertExpandOffsets([1, 2, 3, 4, 5, 6], index, [2])
 
1440
        self.assertExpandOffsets([94, 95, 96, 97, 98, 99], index, [98])
 
1441
 
 
1442
        # Requesting many nodes will expand all locations equally
 
1443
        self.assertExpandOffsets([1, 2, 3, 80, 81, 82], index, [2, 81])
 
1444
        self.assertExpandOffsets([1, 2, 3, 9, 10, 11, 80, 81, 82], index,
 
1445
                               [2, 10, 81])
 
1446
 
 
1447
    def test_stop_at_cached(self):
 
1448
        index = self.make_100_node_index()
 
1449
        self.set_cached_offsets(index, [0, 10, 19])
 
1450
        self.assertExpandOffsets([11, 12, 13, 14, 15, 16], index, [11])
 
1451
        self.assertExpandOffsets([11, 12, 13, 14, 15, 16], index, [12])
 
1452
        self.assertExpandOffsets([12, 13, 14, 15, 16, 17, 18], index, [15])
 
1453
        self.assertExpandOffsets([13, 14, 15, 16, 17, 18], index, [16])
 
1454
        self.assertExpandOffsets([13, 14, 15, 16, 17, 18], index, [17])
 
1455
        self.assertExpandOffsets([13, 14, 15, 16, 17, 18], index, [18])
 
1456
 
 
1457
    def test_cannot_fully_expand(self):
 
1458
        index = self.make_100_node_index()
 
1459
        self.set_cached_offsets(index, [0, 10, 12])
 
1460
        # We don't go into an endless loop if we are bound by cached nodes
 
1461
        self.assertExpandOffsets([11], index, [11])
 
1462
 
 
1463
    def test_overlap(self):
 
1464
        index = self.make_100_node_index()
 
1465
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15], index, [12, 13])
 
1466
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15], index, [11, 14])
 
1467
 
 
1468
    def test_stay_within_layer(self):
 
1469
        index = self.make_1000_node_index()
 
1470
        # When expanding a request, we won't read nodes from the next layer
 
1471
        self.assertExpandOffsets([1, 2, 3, 4], index, [2])
 
1472
        self.assertExpandOffsets([6, 7, 8, 9], index, [6])
 
1473
        self.assertExpandOffsets([6, 7, 8, 9], index, [9])
 
1474
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15], index, [10])
 
1475
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15, 16], index, [13])
 
1476
 
 
1477
        self.set_cached_offsets(index, [0, 4, 12])
 
1478
        self.assertExpandOffsets([5, 6, 7, 8, 9], index, [7])
 
1479
        self.assertExpandOffsets([10, 11], index, [11])
 
1480
 
 
1481
    def test_small_requests_unexpanded(self):
 
1482
        index = self.make_100_node_index()
 
1483
        self.set_cached_offsets(index, [0])
 
1484
        self.assertExpandOffsets([1], index, [1])
 
1485
        self.assertExpandOffsets([50], index, [50])
 
1486
        # If we request more than one node, then we'll expand
 
1487
        self.assertExpandOffsets([49, 50, 51, 59, 60, 61], index, [50, 60])
 
1488
 
 
1489
        # The first pass does not expand
 
1490
        index = self.make_1000_node_index()
 
1491
        self.set_cached_offsets(index, [0])
 
1492
        self.assertExpandOffsets([1], index, [1])
 
1493
        self.set_cached_offsets(index, [0, 1])
 
1494
        self.assertExpandOffsets([100], index, [100])
 
1495
        self.set_cached_offsets(index, [0, 1, 100])
 
1496
        # But after the first depth, we will expand
 
1497
        self.assertExpandOffsets([2, 3, 4, 5, 6, 7], index, [2])
 
1498
        self.assertExpandOffsets([2, 3, 4, 5, 6, 7], index, [4])
 
1499
        self.set_cached_offsets(index, [0, 1, 2, 3, 4, 5, 6, 7, 100])
 
1500
        self.assertExpandOffsets([102, 103, 104, 105, 106, 107, 108], index,
 
1501
                                 [105])