/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: 2008-10-08 21:56:12 UTC
  • mto: This revision was merged to the branch mainline in revision 3773.
  • Revision ID: john@arbash-meinel.com-20081008215612-y9v94tqxreqoangx
Simplify the --raw mode.

I didn't realize, but the only node that is special cased is the 'root' node,
and to read it, you actually have to parse it directly, because the
compressed bytes start immediately after the end of the header, rather than
having any padding before the zlib bytes.

Show diffs side-by-side

added added

removed removed

Lines of Context:
bzrlib/tests/test_btree_index.py
1
 
# Copyright (C) 2008, 2009, 2010 Canonical Ltd
 
1
# Copyright (C) 2008 Canonical Ltd
2
2
#
3
3
# This program is free software; you can redistribute it and/or modify
4
4
# it under the terms of the GNU General Public License as published by
12
12
#
13
13
# You should have received a copy of the GNU General Public License
14
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
 
15
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
16
16
#
17
17
 
18
18
"""Tests for btree indices."""
23
23
from bzrlib import (
24
24
    btree_index,
25
25
    errors,
26
 
    fifo_cache,
27
 
    lru_cache,
28
 
    osutils,
29
26
    tests,
30
27
    )
31
28
from bzrlib.tests import (
32
29
    TestCaseWithTransport,
 
30
    TestScenarioApplier,
 
31
    adapt_tests,
33
32
    condition_isinstance,
34
 
    multiply_tests,
35
33
    split_suite_by_condition,
36
34
    )
37
35
from bzrlib.transport import get_transport
41
39
    # parameterise the TestBTreeNodes tests
42
40
    node_tests, others = split_suite_by_condition(standard_tests,
43
41
        condition_isinstance(TestBTreeNodes))
 
42
    applier = TestScenarioApplier()
44
43
    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')
 
44
    applier.scenarios = [('python', {'parse_btree': py_module})]
 
45
    if CompiledBtreeParserFeature.available():
 
46
        # Is there a way to do this that gets missing feature failures rather
 
47
        # than no indication to the user?
 
48
        import bzrlib._btree_serializer_c as c_module
 
49
        applier.scenarios.append(('C', {'parse_btree': c_module}))
 
50
    adapt_tests(node_tests, applier, others)
 
51
    return others
 
52
 
 
53
 
 
54
class _CompiledBtreeParserFeature(tests.Feature):
 
55
    def _probe(self):
 
56
        try:
 
57
            import bzrlib._btree_serializer_c
 
58
        except ImportError:
 
59
            return False
 
60
        return True
 
61
 
 
62
    def feature_name(self):
 
63
        return 'bzrlib._btree_serializer_c'
 
64
 
 
65
CompiledBtreeParserFeature = _CompiledBtreeParserFeature()
54
66
 
55
67
 
56
68
class BTreeTestCase(TestCaseWithTransport):
59
71
 
60
72
    def setUp(self):
61
73
        TestCaseWithTransport.setUp(self)
62
 
        self.overrideAttr(btree_index, '_RESERVED_HEADER_BYTES', 100)
 
74
        self._original_header = btree_index._RESERVED_HEADER_BYTES
 
75
        def restore():
 
76
            btree_index._RESERVED_HEADER_BYTES = self._original_header
 
77
        self.addCleanup(restore)
 
78
        btree_index._RESERVED_HEADER_BYTES = 100
63
79
 
64
80
    def make_nodes(self, count, key_elements, reference_lists):
65
81
        """Generate count*key_elements sample nodes."""
99
115
 
100
116
    def shrink_page_size(self):
101
117
        """Shrink the default page size so that less fits in a page."""
102
 
        self.overrideAttr(btree_index, '_PAGE_SIZE')
 
118
        old_page_size = btree_index._PAGE_SIZE
 
119
        def cleanup():
 
120
            btree_index._PAGE_SIZE = old_page_size
 
121
        self.addCleanup(cleanup)
103
122
        btree_index._PAGE_SIZE = 2048
104
123
 
105
124
 
106
125
class TestBTreeBuilder(BTreeTestCase):
107
126
 
108
 
    def test_clear_cache(self):
109
 
        builder = btree_index.BTreeBuilder(reference_lists=0, key_elements=1)
110
 
        # This is a no-op, but we need the api to be consistent with other
111
 
        # BTreeGraphIndex apis.
112
 
        builder.clear_cache()
113
 
 
114
127
    def test_empty_1_0(self):
115
128
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
116
129
        # NamedTemporaryFile dies on builder.finish().read(). weird.
142
155
        temp_file = builder.finish()
143
156
        content = temp_file.read()
144
157
        del temp_file
145
 
        self.assertEqual(131, len(content))
 
158
        self.assertEqual(158, len(content))
146
159
        self.assertEqual(
147
160
            "B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=5\n"
148
161
            "row_lengths=1\n",
166
179
        temp_file = builder.finish()
167
180
        content = temp_file.read()
168
181
        del temp_file
169
 
        self.assertEqual(238, len(content))
 
182
        self.assertEqual(264, len(content))
170
183
        self.assertEqual(
171
184
            "B+Tree Graph Index 2\nnode_ref_lists=2\nkey_elements=2\nlen=10\n"
172
185
            "row_lengths=1\n",
232
245
        temp_file = builder.finish()
233
246
        content = temp_file.read()
234
247
        del temp_file
235
 
        self.assertEqual(155, len(content))
 
248
        self.assertEqual(181, len(content))
236
249
        self.assertEqual(
237
250
            "B+Tree Graph Index 2\nnode_ref_lists=0\nkey_elements=1\nlen=10\n"
238
251
            "row_lengths=1\n",
340
353
        # Test the parts of the index that take up memory are doing so
341
354
        # predictably.
342
355
        self.assertEqual(1, len(builder._nodes))
 
356
        self.assertEqual(1, len(builder._keys))
343
357
        self.assertIs(None, builder._nodes_by_key)
344
358
        builder.add_node(*nodes[1])
345
359
        self.assertEqual(0, len(builder._nodes))
 
360
        self.assertEqual(0, len(builder._keys))
346
361
        self.assertIs(None, builder._nodes_by_key)
347
362
        self.assertEqual(1, len(builder._backing_indices))
348
363
        self.assertEqual(2, builder._backing_indices[0].key_count())
349
364
        # now back to memory
350
365
        builder.add_node(*nodes[2])
351
366
        self.assertEqual(1, len(builder._nodes))
 
367
        self.assertEqual(1, len(builder._keys))
352
368
        self.assertIs(None, builder._nodes_by_key)
353
369
        # And spills to a second backing index combing all
354
370
        builder.add_node(*nodes[3])
355
371
        self.assertEqual(0, len(builder._nodes))
 
372
        self.assertEqual(0, len(builder._keys))
356
373
        self.assertIs(None, builder._nodes_by_key)
357
374
        self.assertEqual(2, len(builder._backing_indices))
358
375
        self.assertEqual(None, builder._backing_indices[0])
361
378
        builder.add_node(*nodes[4])
362
379
        builder.add_node(*nodes[5])
363
380
        self.assertEqual(0, len(builder._nodes))
 
381
        self.assertEqual(0, len(builder._keys))
364
382
        self.assertIs(None, builder._nodes_by_key)
365
383
        self.assertEqual(2, len(builder._backing_indices))
366
384
        self.assertEqual(2, builder._backing_indices[0].key_count())
416
434
        self.assertEqual(sorted(nodes), nodes)
417
435
        self.assertEqual(16, len(nodes))
418
436
 
419
 
    def test_spill_index_stress_1_1_no_combine(self):
420
 
        builder = btree_index.BTreeBuilder(key_elements=1, spill_at=2)
421
 
        builder.set_optimize(for_size=False, combine_backing_indices=False)
422
 
        nodes = [node[0:2] for node in self.make_nodes(16, 1, 0)]
423
 
        builder.add_node(*nodes[0])
424
 
        # Test the parts of the index that take up memory are doing so
425
 
        # predictably.
426
 
        self.assertEqual(1, len(builder._nodes))
427
 
        self.assertIs(None, builder._nodes_by_key)
428
 
        builder.add_node(*nodes[1])
429
 
        self.assertEqual(0, len(builder._nodes))
430
 
        self.assertIs(None, builder._nodes_by_key)
431
 
        self.assertEqual(1, len(builder._backing_indices))
432
 
        self.assertEqual(2, builder._backing_indices[0].key_count())
433
 
        # now back to memory
434
 
        builder.add_node(*nodes[2])
435
 
        self.assertEqual(1, len(builder._nodes))
436
 
        self.assertIs(None, builder._nodes_by_key)
437
 
        # And spills to a second backing index but doesn't combine
438
 
        builder.add_node(*nodes[3])
439
 
        self.assertEqual(0, len(builder._nodes))
440
 
        self.assertIs(None, builder._nodes_by_key)
441
 
        self.assertEqual(2, len(builder._backing_indices))
442
 
        for backing_index in builder._backing_indices:
443
 
            self.assertEqual(2, backing_index.key_count())
444
 
        # The next spills to the 3rd slot
445
 
        builder.add_node(*nodes[4])
446
 
        builder.add_node(*nodes[5])
447
 
        self.assertEqual(0, len(builder._nodes))
448
 
        self.assertIs(None, builder._nodes_by_key)
449
 
        self.assertEqual(3, len(builder._backing_indices))
450
 
        for backing_index in builder._backing_indices:
451
 
            self.assertEqual(2, backing_index.key_count())
452
 
        # Now spill a few more, and check that we don't combine
453
 
        builder.add_node(*nodes[6])
454
 
        builder.add_node(*nodes[7])
455
 
        builder.add_node(*nodes[8])
456
 
        builder.add_node(*nodes[9])
457
 
        builder.add_node(*nodes[10])
458
 
        builder.add_node(*nodes[11])
459
 
        builder.add_node(*nodes[12])
460
 
        self.assertEqual(6, len(builder._backing_indices))
461
 
        for backing_index in builder._backing_indices:
462
 
            self.assertEqual(2, backing_index.key_count())
463
 
        # Test that memory and disk are both used for query methods; and that
464
 
        # None is skipped over happily.
465
 
        self.assertEqual([(builder,) + node for node in sorted(nodes[:13])],
466
 
            list(builder.iter_all_entries()))
467
 
        # Two nodes - one memory one disk
468
 
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
469
 
            set(builder.iter_entries([nodes[12][0], nodes[11][0]])))
470
 
        self.assertEqual(13, builder.key_count())
471
 
        self.assertEqual(set([(builder,) + node for node in nodes[11:13]]),
472
 
            set(builder.iter_entries_prefix([nodes[12][0], nodes[11][0]])))
473
 
        builder.add_node(*nodes[13])
474
 
        builder.add_node(*nodes[14])
475
 
        builder.add_node(*nodes[15])
476
 
        self.assertEqual(8, len(builder._backing_indices))
477
 
        for backing_index in builder._backing_indices:
478
 
            self.assertEqual(2, backing_index.key_count())
479
 
        # Now finish, and check we got a correctly ordered tree
480
 
        transport = self.get_transport('')
481
 
        size = transport.put_file('index', builder.finish())
482
 
        index = btree_index.BTreeGraphIndex(transport, 'index', size)
483
 
        nodes = list(index.iter_all_entries())
484
 
        self.assertEqual(sorted(nodes), nodes)
485
 
        self.assertEqual(16, len(nodes))
486
 
 
487
 
    def test_set_optimize(self):
488
 
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
489
 
        builder.set_optimize(for_size=True)
490
 
        self.assertTrue(builder._optimize_for_size)
491
 
        builder.set_optimize(for_size=False)
492
 
        self.assertFalse(builder._optimize_for_size)
493
 
        # test that we can set combine_backing_indices without effecting
494
 
        # _optimize_for_size
495
 
        obj = object()
496
 
        builder._optimize_for_size = obj
497
 
        builder.set_optimize(combine_backing_indices=False)
498
 
        self.assertFalse(builder._combine_backing_indices)
499
 
        self.assertIs(obj, builder._optimize_for_size)
500
 
        builder.set_optimize(combine_backing_indices=True)
501
 
        self.assertTrue(builder._combine_backing_indices)
502
 
        self.assertIs(obj, builder._optimize_for_size)
503
 
 
504
437
    def test_spill_index_stress_2_2(self):
505
438
        # test that references and longer keys don't confuse things.
506
439
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2,
509
442
        builder.add_node(*nodes[0])
510
443
        # Test the parts of the index that take up memory are doing so
511
444
        # predictably.
 
445
        self.assertEqual(1, len(builder._keys))
512
446
        self.assertEqual(1, len(builder._nodes))
513
447
        self.assertIs(None, builder._nodes_by_key)
514
448
        builder.add_node(*nodes[1])
 
449
        self.assertEqual(0, len(builder._keys))
515
450
        self.assertEqual(0, len(builder._nodes))
516
451
        self.assertIs(None, builder._nodes_by_key)
517
452
        self.assertEqual(1, len(builder._backing_indices))
520
455
        old = dict(builder._get_nodes_by_key()) #Build up the nodes by key dict
521
456
        builder.add_node(*nodes[2])
522
457
        self.assertEqual(1, len(builder._nodes))
 
458
        self.assertEqual(1, len(builder._keys))
523
459
        self.assertIsNot(None, builder._nodes_by_key)
524
460
        self.assertNotEqual({}, builder._nodes_by_key)
525
461
        # We should have a new entry
527
463
        # And spills to a second backing index combing all
528
464
        builder.add_node(*nodes[3])
529
465
        self.assertEqual(0, len(builder._nodes))
 
466
        self.assertEqual(0, len(builder._keys))
530
467
        self.assertIs(None, builder._nodes_by_key)
531
468
        self.assertEqual(2, len(builder._backing_indices))
532
469
        self.assertEqual(None, builder._backing_indices[0])
535
472
        builder.add_node(*nodes[4])
536
473
        builder.add_node(*nodes[5])
537
474
        self.assertEqual(0, len(builder._nodes))
 
475
        self.assertEqual(0, len(builder._keys))
538
476
        self.assertIs(None, builder._nodes_by_key)
539
477
        self.assertEqual(2, len(builder._backing_indices))
540
478
        self.assertEqual(2, builder._backing_indices[0].key_count())
611
549
        size = trans.put_file('index', stream)
612
550
        return btree_index.BTreeGraphIndex(trans, 'index', size)
613
551
 
614
 
    def make_index_with_offset(self, ref_lists=1, key_elements=1, nodes=[],
615
 
                               offset=0):
616
 
        builder = btree_index.BTreeBuilder(key_elements=key_elements,
617
 
                                           reference_lists=ref_lists)
618
 
        builder.add_nodes(nodes)
619
 
        transport = self.get_transport('')
620
 
        # NamedTemporaryFile dies on builder.finish().read(). weird.
621
 
        temp_file = builder.finish()
622
 
        content = temp_file.read()
623
 
        del temp_file
624
 
        size = len(content)
625
 
        transport.put_bytes('index', (' '*offset)+content)
626
 
        return btree_index.BTreeGraphIndex(transport, 'index', size=size,
627
 
                                           offset=offset)
628
 
 
629
 
    def test_clear_cache(self):
630
 
        nodes = self.make_nodes(160, 2, 2)
631
 
        index = self.make_index(ref_lists=2, key_elements=2, nodes=nodes)
632
 
        self.assertEqual(1, len(list(index.iter_entries([nodes[30][0]]))))
633
 
        self.assertEqual([1, 4], index._row_lengths)
634
 
        self.assertIsNot(None, index._root_node)
635
 
        internal_node_pre_clear = index._internal_node_cache.keys()
636
 
        self.assertTrue(len(index._leaf_node_cache) > 0)
637
 
        index.clear_cache()
638
 
        # We don't touch _root_node or _internal_node_cache, both should be
639
 
        # small, and can save a round trip or two
640
 
        self.assertIsNot(None, index._root_node)
641
 
        # NOTE: We don't want to affect the _internal_node_cache, as we expect
642
 
        #       it will be small, and if we ever do touch this index again, it
643
 
        #       will save round-trips.  This assertion isn't very strong,
644
 
        #       becuase without a 3-level index, we don't have any internal
645
 
        #       nodes cached.
646
 
        self.assertEqual(internal_node_pre_clear,
647
 
                         index._internal_node_cache.keys())
648
 
        self.assertEqual(0, len(index._leaf_node_cache))
649
 
 
650
552
    def test_trivial_constructor(self):
651
553
        transport = get_transport('trace+' + self.get_url(''))
652
554
        index = btree_index.BTreeGraphIndex(transport, 'index', None)
670
572
        # The entire index should have been requested (as we generally have the
671
573
        # size available, and doing many small readvs is inappropriate).
672
574
        # We can't tell how much was actually read here, but - check the code.
673
 
        self.assertEqual([('get', 'index')], transport._activity)
 
575
        self.assertEqual([('get', 'index'),
 
576
            ('readv', 'index', [(0, 72)], False, None)],
 
577
            transport._activity)
674
578
 
675
579
    def test_empty_key_count(self):
676
580
        builder = btree_index.BTreeBuilder(key_elements=1, reference_lists=0)
699
603
        # The entire index should have been read, as it is one page long.
700
604
        self.assertEqual([('readv', 'index', [(0, size)], False, None)],
701
605
            transport._activity)
702
 
        self.assertEqual(1173, size)
703
 
 
704
 
    def test_with_offset_no_size(self):
705
 
        index = self.make_index_with_offset(key_elements=1, ref_lists=1,
706
 
                                            offset=1234,
707
 
                                            nodes=self.make_nodes(200, 1, 1))
708
 
        index._size = None # throw away the size info
709
 
        self.assertEqual(200, index.key_count())
710
 
 
711
 
    def test_with_small_offset(self):
712
 
        index = self.make_index_with_offset(key_elements=1, ref_lists=1,
713
 
                                            offset=1234,
714
 
                                            nodes=self.make_nodes(200, 1, 1))
715
 
        self.assertEqual(200, index.key_count())
716
 
 
717
 
    def test_with_large_offset(self):
718
 
        index = self.make_index_with_offset(key_elements=1, ref_lists=1,
719
 
                                            offset=123456,
720
 
                                            nodes=self.make_nodes(200, 1, 1))
721
 
        self.assertEqual(200, index.key_count())
722
 
 
723
 
    def test__read_nodes_no_size_one_page_reads_once(self):
724
 
        self.make_index(nodes=[(('key',), 'value', ())])
725
 
        trans = get_transport('trace+' + self.get_url())
726
 
        index = btree_index.BTreeGraphIndex(trans, 'index', None)
727
 
        del trans._activity[:]
728
 
        nodes = dict(index._read_nodes([0]))
729
 
        self.assertEqual([0], nodes.keys())
730
 
        node = nodes[0]
731
 
        self.assertEqual([('key',)], node.keys.keys())
732
 
        self.assertEqual([('get', 'index')], trans._activity)
733
 
 
734
 
    def test__read_nodes_no_size_multiple_pages(self):
735
 
        index = self.make_index(2, 2, nodes=self.make_nodes(160, 2, 2))
736
 
        index.key_count()
737
 
        num_pages = index._row_offsets[-1]
738
 
        # Reopen with a traced transport and no size
739
 
        trans = get_transport('trace+' + self.get_url())
740
 
        index = btree_index.BTreeGraphIndex(trans, 'index', None)
741
 
        del trans._activity[:]
742
 
        nodes = dict(index._read_nodes([0]))
743
 
        self.assertEqual(range(num_pages), nodes.keys())
 
606
        self.assertEqual(1199, size)
744
607
 
745
608
    def test_2_levels_key_count_2_2(self):
746
609
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
772
635
        # The entire index should have been read linearly.
773
636
        self.assertEqual([('readv', 'index', [(0, size)], False, None)],
774
637
            transport._activity)
775
 
        self.assertEqual(1488, size)
 
638
        self.assertEqual(1514, size)
776
639
 
777
640
    def test_validate_two_pages(self):
778
641
        builder = btree_index.BTreeBuilder(key_elements=2, reference_lists=2)
829
692
            btree_index.BTreeGraphIndex(transport1, 'index', 10) !=
830
693
            btree_index.BTreeGraphIndex(transport1, 'index', 20))
831
694
 
832
 
    def test_iter_all_only_root_no_size(self):
833
 
        self.make_index(nodes=[(('key',), 'value', ())])
834
 
        trans = get_transport('trace+' + self.get_url(''))
835
 
        index = btree_index.BTreeGraphIndex(trans, 'index', None)
836
 
        del trans._activity[:]
837
 
        self.assertEqual([(('key',), 'value')],
838
 
                         [x[1:] for x in index.iter_all_entries()])
839
 
        self.assertEqual([('get', 'index')], trans._activity)
840
 
 
841
695
    def test_iter_all_entries_reads(self):
842
696
        # iterating all entries reads the header, then does a linear
843
697
        # read.
979
833
            (index, ('name', 'fin2'), 'beta', ((), ))]),
980
834
            set(index.iter_entries_prefix([('name', None)])))
981
835
 
982
 
    # XXX: external_references tests are duplicated in test_index.  We
983
 
    # probably should have per_graph_index tests...
984
 
    def test_external_references_no_refs(self):
985
 
        index = self.make_index(ref_lists=0, nodes=[])
986
 
        self.assertRaises(ValueError, index.external_references, 0)
987
 
 
988
 
    def test_external_references_no_results(self):
989
 
        index = self.make_index(ref_lists=1, nodes=[
990
 
            (('key',), 'value', ([],))])
991
 
        self.assertEqual(set(), index.external_references(0))
992
 
 
993
 
    def test_external_references_missing_ref(self):
994
 
        missing_key = ('missing',)
995
 
        index = self.make_index(ref_lists=1, nodes=[
996
 
            (('key',), 'value', ([missing_key],))])
997
 
        self.assertEqual(set([missing_key]), index.external_references(0))
998
 
 
999
 
    def test_external_references_multiple_ref_lists(self):
1000
 
        missing_key = ('missing',)
1001
 
        index = self.make_index(ref_lists=2, nodes=[
1002
 
            (('key',), 'value', ([], [missing_key]))])
1003
 
        self.assertEqual(set([]), index.external_references(0))
1004
 
        self.assertEqual(set([missing_key]), index.external_references(1))
1005
 
 
1006
 
    def test_external_references_two_records(self):
1007
 
        index = self.make_index(ref_lists=1, nodes=[
1008
 
            (('key-1',), 'value', ([('key-2',)],)),
1009
 
            (('key-2',), 'value', ([],)),
1010
 
            ])
1011
 
        self.assertEqual(set([]), index.external_references(0))
1012
 
 
1013
 
    def test__find_ancestors_one_page(self):
1014
 
        key1 = ('key-1',)
1015
 
        key2 = ('key-2',)
1016
 
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
1017
 
            (key1, 'value', ([key2],)),
1018
 
            (key2, 'value', ([],)),
1019
 
            ])
1020
 
        parent_map = {}
1021
 
        missing_keys = set()
1022
 
        search_keys = index._find_ancestors([key1], 0, parent_map, missing_keys)
1023
 
        self.assertEqual({key1: (key2,), key2: ()}, parent_map)
1024
 
        self.assertEqual(set(), missing_keys)
1025
 
        self.assertEqual(set(), search_keys)
1026
 
 
1027
 
    def test__find_ancestors_one_page_w_missing(self):
1028
 
        key1 = ('key-1',)
1029
 
        key2 = ('key-2',)
1030
 
        key3 = ('key-3',)
1031
 
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
1032
 
            (key1, 'value', ([key2],)),
1033
 
            (key2, 'value', ([],)),
1034
 
            ])
1035
 
        parent_map = {}
1036
 
        missing_keys = set()
1037
 
        search_keys = index._find_ancestors([key2, key3], 0, parent_map,
1038
 
                                            missing_keys)
1039
 
        self.assertEqual({key2: ()}, parent_map)
1040
 
        # we know that key3 is missing because we read the page that it would
1041
 
        # otherwise be on
1042
 
        self.assertEqual(set([key3]), missing_keys)
1043
 
        self.assertEqual(set(), search_keys)
1044
 
 
1045
 
    def test__find_ancestors_one_parent_missing(self):
1046
 
        key1 = ('key-1',)
1047
 
        key2 = ('key-2',)
1048
 
        key3 = ('key-3',)
1049
 
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
1050
 
            (key1, 'value', ([key2],)),
1051
 
            (key2, 'value', ([key3],)),
1052
 
            ])
1053
 
        parent_map = {}
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
 
        # all we know is that key3 wasn't present on the page we were reading
1060
 
        # but if you look, the last key is key2 which comes before key3, so we
1061
 
        # don't know whether key3 would land on this page or not.
1062
 
        self.assertEqual(set([key3]), search_keys)
1063
 
        search_keys = index._find_ancestors(search_keys, 0, parent_map,
1064
 
                                            missing_keys)
1065
 
        # passing it back in, we are sure it is 'missing'
1066
 
        self.assertEqual({key1: (key2,), key2: (key3,)}, parent_map)
1067
 
        self.assertEqual(set([key3]), missing_keys)
1068
 
        self.assertEqual(set([]), search_keys)
1069
 
 
1070
 
    def test__find_ancestors_dont_search_known(self):
1071
 
        key1 = ('key-1',)
1072
 
        key2 = ('key-2',)
1073
 
        key3 = ('key-3',)
1074
 
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[
1075
 
            (key1, 'value', ([key2],)),
1076
 
            (key2, 'value', ([key3],)),
1077
 
            (key3, 'value', ([],)),
1078
 
            ])
1079
 
        # We already know about key2, so we won't try to search for key3
1080
 
        parent_map = {key2: (key3,)}
1081
 
        missing_keys = set()
1082
 
        search_keys = index._find_ancestors([key1], 0, parent_map,
1083
 
                                            missing_keys)
1084
 
        self.assertEqual({key1: (key2,), key2: (key3,)}, parent_map)
1085
 
        self.assertEqual(set(), missing_keys)
1086
 
        self.assertEqual(set(), search_keys)
1087
 
 
1088
 
    def test__find_ancestors_multiple_pages(self):
1089
 
        # We need to use enough keys that we actually cause a split
1090
 
        start_time = 1249671539
1091
 
        email = "joebob@example.com"
1092
 
        nodes = []
1093
 
        ref_lists = ((),)
1094
 
        rev_keys = []
1095
 
        for i in xrange(400):
1096
 
            rev_id = '%s-%s-%s' % (email,
1097
 
                                   osutils.compact_date(start_time + i),
1098
 
                                   osutils.rand_chars(16))
1099
 
            rev_key = (rev_id,)
1100
 
            nodes.append((rev_key, 'value', ref_lists))
1101
 
            # We have a ref 'list' of length 1, with a list of parents, with 1
1102
 
            # parent which is a key
1103
 
            ref_lists = ((rev_key,),)
1104
 
            rev_keys.append(rev_key)
1105
 
        index = self.make_index(ref_lists=1, key_elements=1, nodes=nodes)
1106
 
        self.assertEqual(400, index.key_count())
1107
 
        self.assertEqual(3, len(index._row_offsets))
1108
 
        nodes = dict(index._read_nodes([1, 2]))
1109
 
        l1 = nodes[1]
1110
 
        l2 = nodes[2]
1111
 
        min_l2_key = l2.min_key
1112
 
        max_l1_key = l1.max_key
1113
 
        self.assertTrue(max_l1_key < min_l2_key)
1114
 
        parents_min_l2_key = l2.keys[min_l2_key][1][0]
1115
 
        self.assertEqual((l1.max_key,), parents_min_l2_key)
1116
 
        # Now, whatever key we select that would fall on the second page,
1117
 
        # should give us all the parents until the page break
1118
 
        key_idx = rev_keys.index(min_l2_key)
1119
 
        next_key = rev_keys[key_idx+1]
1120
 
        # So now when we get the parent map, we should get the key we are
1121
 
        # looking for, min_l2_key, and then a reference to go look for the
1122
 
        # parent of that key
1123
 
        parent_map = {}
1124
 
        missing_keys = set()
1125
 
        search_keys = index._find_ancestors([next_key], 0, parent_map,
1126
 
                                            missing_keys)
1127
 
        self.assertEqual([min_l2_key, next_key], sorted(parent_map))
1128
 
        self.assertEqual(set(), missing_keys)
1129
 
        self.assertEqual(set([max_l1_key]), search_keys)
1130
 
        parent_map = {}
1131
 
        search_keys = index._find_ancestors([max_l1_key], 0, parent_map,
1132
 
                                            missing_keys)
1133
 
        self.assertEqual(sorted(l1.keys), sorted(parent_map))
1134
 
        self.assertEqual(set(), missing_keys)
1135
 
        self.assertEqual(set(), search_keys)
1136
 
 
1137
 
    def test__find_ancestors_empty_index(self):
1138
 
        index = self.make_index(ref_lists=1, key_elements=1, nodes=[])
1139
 
        parent_map = {}
1140
 
        missing_keys = set()
1141
 
        search_keys = index._find_ancestors([('one',), ('two',)], 0, parent_map,
1142
 
                                            missing_keys)
1143
 
        self.assertEqual(set(), search_keys)
1144
 
        self.assertEqual({}, parent_map)
1145
 
        self.assertEqual(set([('one',), ('two',)]), missing_keys)
1146
 
 
1147
 
    def test_supports_unlimited_cache(self):
1148
 
        builder = btree_index.BTreeBuilder(reference_lists=0, key_elements=1)
1149
 
        # We need enough nodes to cause a page split (so we have both an
1150
 
        # internal node and a couple leaf nodes. 500 seems to be enough.)
1151
 
        nodes = self.make_nodes(500, 1, 0)
1152
 
        for node in nodes:
1153
 
            builder.add_node(*node)
1154
 
        stream = builder.finish()
1155
 
        trans = get_transport(self.get_url())
1156
 
        size = trans.put_file('index', stream)
1157
 
        index = btree_index.BTreeGraphIndex(trans, 'index', size)
1158
 
        self.assertEqual(500, index.key_count())
1159
 
        # We have an internal node
1160
 
        self.assertEqual(2, len(index._row_lengths))
1161
 
        # We have at least 2 leaf nodes
1162
 
        self.assertTrue(index._row_lengths[-1] >= 2)
1163
 
        self.assertIsInstance(index._leaf_node_cache, lru_cache.LRUCache)
1164
 
        self.assertEqual(btree_index._NODE_CACHE_SIZE,
1165
 
                         index._leaf_node_cache._max_cache)
1166
 
        self.assertIsInstance(index._internal_node_cache, fifo_cache.FIFOCache)
1167
 
        self.assertEqual(100, index._internal_node_cache._max_cache)
1168
 
        # No change if unlimited_cache=False is passed
1169
 
        index = btree_index.BTreeGraphIndex(trans, 'index', size,
1170
 
                                            unlimited_cache=False)
1171
 
        self.assertIsInstance(index._leaf_node_cache, lru_cache.LRUCache)
1172
 
        self.assertEqual(btree_index._NODE_CACHE_SIZE,
1173
 
                         index._leaf_node_cache._max_cache)
1174
 
        self.assertIsInstance(index._internal_node_cache, fifo_cache.FIFOCache)
1175
 
        self.assertEqual(100, index._internal_node_cache._max_cache)
1176
 
        index = btree_index.BTreeGraphIndex(trans, 'index', size,
1177
 
                                            unlimited_cache=True)
1178
 
        self.assertIsInstance(index._leaf_node_cache, dict)
1179
 
        self.assertIs(type(index._internal_node_cache), dict)
1180
 
        # Exercise the lookup code
1181
 
        entries = set(index.iter_entries([n[0] for n in nodes]))
1182
 
        self.assertEqual(500, len(entries))
1183
 
 
1184
836
 
1185
837
class TestBTreeNodes(BTreeTestCase):
1186
838
 
 
839
    def restore_parser(self):
 
840
        btree_index._btree_serializer = self.saved_parser
 
841
 
1187
842
    def setUp(self):
1188
843
        BTreeTestCase.setUp(self)
1189
 
        self.overrideAttr(btree_index, '_btree_serializer', self.parse_btree)
 
844
        self.saved_parser = btree_index._btree_serializer
 
845
        self.addCleanup(self.restore_parser)
 
846
        btree_index._btree_serializer = self.parse_btree
1190
847
 
1191
848
    def test_LeafNode_1_0(self):
1192
849
        node_bytes = ("type=leaf\n"
1303
960
    def test_exists(self):
1304
961
        # This is just to let the user know if they don't have the feature
1305
962
        # available
1306
 
        self.requireFeature(compiled_btreeparser_feature)
 
963
        self.requireFeature(CompiledBtreeParserFeature)
1307
964
 
1308
965
 
1309
966
class TestMultiBisectRight(tests.TestCase):
1339
996
                                     (4, ['g', 'h'])],
1340
997
                                    ['a', 'b', 'd', 'e', 'g', 'h'],
1341
998
                                    ['c', 'd', 'f', 'g'])
1342
 
 
1343
 
 
1344
 
class TestExpandOffsets(tests.TestCase):
1345
 
 
1346
 
    def make_index(self, size, recommended_pages=None):
1347
 
        """Make an index with a generic size.
1348
 
 
1349
 
        This doesn't actually create anything on disk, it just primes a
1350
 
        BTreeGraphIndex with the recommended information.
1351
 
        """
1352
 
        index = btree_index.BTreeGraphIndex(get_transport('memory:///'),
1353
 
                                            'test-index', size=size)
1354
 
        if recommended_pages is not None:
1355
 
            index._recommended_pages = recommended_pages
1356
 
        return index
1357
 
 
1358
 
    def set_cached_offsets(self, index, cached_offsets):
1359
 
        """Monkeypatch to give a canned answer for _get_offsets_for...()."""
1360
 
        def _get_offsets_to_cached_pages():
1361
 
            cached = set(cached_offsets)
1362
 
            return cached
1363
 
        index._get_offsets_to_cached_pages = _get_offsets_to_cached_pages
1364
 
 
1365
 
    def prepare_index(self, index, node_ref_lists, key_length, key_count,
1366
 
                      row_lengths, cached_offsets):
1367
 
        """Setup the BTreeGraphIndex with some pre-canned information."""
1368
 
        index.node_ref_lists = node_ref_lists
1369
 
        index._key_length = key_length
1370
 
        index._key_count = key_count
1371
 
        index._row_lengths = row_lengths
1372
 
        index._compute_row_offsets()
1373
 
        index._root_node = btree_index._InternalNode('internal\noffset=0\n')
1374
 
        self.set_cached_offsets(index, cached_offsets)
1375
 
 
1376
 
    def make_100_node_index(self):
1377
 
        index = self.make_index(4096*100, 6)
1378
 
        # Consider we've already made a single request at the middle
1379
 
        self.prepare_index(index, node_ref_lists=0, key_length=1,
1380
 
                           key_count=1000, row_lengths=[1, 99],
1381
 
                           cached_offsets=[0, 50])
1382
 
        return index
1383
 
 
1384
 
    def make_1000_node_index(self):
1385
 
        index = self.make_index(4096*1000, 6)
1386
 
        # Pretend we've already made a single request in the middle
1387
 
        self.prepare_index(index, node_ref_lists=0, key_length=1,
1388
 
                           key_count=90000, row_lengths=[1, 9, 990],
1389
 
                           cached_offsets=[0, 5, 500])
1390
 
        return index
1391
 
 
1392
 
    def assertNumPages(self, expected_pages, index, size):
1393
 
        index._size = size
1394
 
        self.assertEqual(expected_pages, index._compute_total_pages_in_index())
1395
 
 
1396
 
    def assertExpandOffsets(self, expected, index, offsets):
1397
 
        self.assertEqual(expected, index._expand_offsets(offsets),
1398
 
                         'We did not get the expected value after expanding'
1399
 
                         ' %s' % (offsets,))
1400
 
 
1401
 
    def test_default_recommended_pages(self):
1402
 
        index = self.make_index(None)
1403
 
        # local transport recommends 4096 byte reads, which is 1 page
1404
 
        self.assertEqual(1, index._recommended_pages)
1405
 
 
1406
 
    def test__compute_total_pages_in_index(self):
1407
 
        index = self.make_index(None)
1408
 
        self.assertNumPages(1, index, 1024)
1409
 
        self.assertNumPages(1, index, 4095)
1410
 
        self.assertNumPages(1, index, 4096)
1411
 
        self.assertNumPages(2, index, 4097)
1412
 
        self.assertNumPages(2, index, 8192)
1413
 
        self.assertNumPages(76, index, 4096*75 + 10)
1414
 
 
1415
 
    def test__find_layer_start_and_stop(self):
1416
 
        index = self.make_1000_node_index()
1417
 
        self.assertEqual((0, 1), index._find_layer_first_and_end(0))
1418
 
        self.assertEqual((1, 10), index._find_layer_first_and_end(1))
1419
 
        self.assertEqual((1, 10), index._find_layer_first_and_end(9))
1420
 
        self.assertEqual((10, 1000), index._find_layer_first_and_end(10))
1421
 
        self.assertEqual((10, 1000), index._find_layer_first_and_end(99))
1422
 
        self.assertEqual((10, 1000), index._find_layer_first_and_end(999))
1423
 
 
1424
 
    def test_unknown_size(self):
1425
 
        # We should not expand if we don't know the file size
1426
 
        index = self.make_index(None, 10)
1427
 
        self.assertExpandOffsets([0], index, [0])
1428
 
        self.assertExpandOffsets([1, 4, 9], index, [1, 4, 9])
1429
 
 
1430
 
    def test_more_than_recommended(self):
1431
 
        index = self.make_index(4096*100, 2)
1432
 
        self.assertExpandOffsets([1, 10], index, [1, 10])
1433
 
        self.assertExpandOffsets([1, 10, 20], index, [1, 10, 20])
1434
 
 
1435
 
    def test_read_all_from_root(self):
1436
 
        index = self.make_index(4096*10, 20)
1437
 
        self.assertExpandOffsets(range(10), index, [0])
1438
 
 
1439
 
    def test_read_all_when_cached(self):
1440
 
        # We've read enough that we can grab all the rest in a single request
1441
 
        index = self.make_index(4096*10, 5)
1442
 
        self.prepare_index(index, node_ref_lists=0, key_length=1,
1443
 
                           key_count=1000, row_lengths=[1, 9],
1444
 
                           cached_offsets=[0, 1, 2, 5, 6])
1445
 
        # It should fill the remaining nodes, regardless of the one requested
1446
 
        self.assertExpandOffsets([3, 4, 7, 8, 9], index, [3])
1447
 
        self.assertExpandOffsets([3, 4, 7, 8, 9], index, [8])
1448
 
        self.assertExpandOffsets([3, 4, 7, 8, 9], index, [9])
1449
 
 
1450
 
    def test_no_root_node(self):
1451
 
        index = self.make_index(4096*10, 5)
1452
 
        self.assertExpandOffsets([0], index, [0])
1453
 
 
1454
 
    def test_include_neighbors(self):
1455
 
        index = self.make_100_node_index()
1456
 
        # We expand in both directions, until we have at least 'recommended'
1457
 
        # pages
1458
 
        self.assertExpandOffsets([9, 10, 11, 12, 13, 14, 15], index, [12])
1459
 
        self.assertExpandOffsets([88, 89, 90, 91, 92, 93, 94], index, [91])
1460
 
        # If we hit an 'edge' we continue in the other direction
1461
 
        self.assertExpandOffsets([1, 2, 3, 4, 5, 6], index, [2])
1462
 
        self.assertExpandOffsets([94, 95, 96, 97, 98, 99], index, [98])
1463
 
 
1464
 
        # Requesting many nodes will expand all locations equally
1465
 
        self.assertExpandOffsets([1, 2, 3, 80, 81, 82], index, [2, 81])
1466
 
        self.assertExpandOffsets([1, 2, 3, 9, 10, 11, 80, 81, 82], index,
1467
 
                               [2, 10, 81])
1468
 
 
1469
 
    def test_stop_at_cached(self):
1470
 
        index = self.make_100_node_index()
1471
 
        self.set_cached_offsets(index, [0, 10, 19])
1472
 
        self.assertExpandOffsets([11, 12, 13, 14, 15, 16], index, [11])
1473
 
        self.assertExpandOffsets([11, 12, 13, 14, 15, 16], index, [12])
1474
 
        self.assertExpandOffsets([12, 13, 14, 15, 16, 17, 18], index, [15])
1475
 
        self.assertExpandOffsets([13, 14, 15, 16, 17, 18], index, [16])
1476
 
        self.assertExpandOffsets([13, 14, 15, 16, 17, 18], index, [17])
1477
 
        self.assertExpandOffsets([13, 14, 15, 16, 17, 18], index, [18])
1478
 
 
1479
 
    def test_cannot_fully_expand(self):
1480
 
        index = self.make_100_node_index()
1481
 
        self.set_cached_offsets(index, [0, 10, 12])
1482
 
        # We don't go into an endless loop if we are bound by cached nodes
1483
 
        self.assertExpandOffsets([11], index, [11])
1484
 
 
1485
 
    def test_overlap(self):
1486
 
        index = self.make_100_node_index()
1487
 
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15], index, [12, 13])
1488
 
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15], index, [11, 14])
1489
 
 
1490
 
    def test_stay_within_layer(self):
1491
 
        index = self.make_1000_node_index()
1492
 
        # When expanding a request, we won't read nodes from the next layer
1493
 
        self.assertExpandOffsets([1, 2, 3, 4], index, [2])
1494
 
        self.assertExpandOffsets([6, 7, 8, 9], index, [6])
1495
 
        self.assertExpandOffsets([6, 7, 8, 9], index, [9])
1496
 
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15], index, [10])
1497
 
        self.assertExpandOffsets([10, 11, 12, 13, 14, 15, 16], index, [13])
1498
 
 
1499
 
        self.set_cached_offsets(index, [0, 4, 12])
1500
 
        self.assertExpandOffsets([5, 6, 7, 8, 9], index, [7])
1501
 
        self.assertExpandOffsets([10, 11], index, [11])
1502
 
 
1503
 
    def test_small_requests_unexpanded(self):
1504
 
        index = self.make_100_node_index()
1505
 
        self.set_cached_offsets(index, [0])
1506
 
        self.assertExpandOffsets([1], index, [1])
1507
 
        self.assertExpandOffsets([50], index, [50])
1508
 
        # If we request more than one node, then we'll expand
1509
 
        self.assertExpandOffsets([49, 50, 51, 59, 60, 61], index, [50, 60])
1510
 
 
1511
 
        # The first pass does not expand
1512
 
        index = self.make_1000_node_index()
1513
 
        self.set_cached_offsets(index, [0])
1514
 
        self.assertExpandOffsets([1], index, [1])
1515
 
        self.set_cached_offsets(index, [0, 1])
1516
 
        self.assertExpandOffsets([100], index, [100])
1517
 
        self.set_cached_offsets(index, [0, 1, 100])
1518
 
        # But after the first depth, we will expand
1519
 
        self.assertExpandOffsets([2, 3, 4, 5, 6, 7], index, [2])
1520
 
        self.assertExpandOffsets([2, 3, 4, 5, 6, 7], index, [4])
1521
 
        self.set_cached_offsets(index, [0, 1, 2, 3, 4, 5, 6, 7, 100])
1522
 
        self.assertExpandOffsets([102, 103, 104, 105, 106, 107, 108], index,
1523
 
                                 [105])