github.com/MetalBlockchain/metalgo@v1.11.9/snow/engine/snowman/ancestor/tree_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package ancestor
     5  
     6  import (
     7  	"testing"
     8  
     9  	"github.com/stretchr/testify/require"
    10  
    11  	"github.com/MetalBlockchain/metalgo/ids"
    12  	"github.com/MetalBlockchain/metalgo/utils/set"
    13  )
    14  
    15  var (
    16  	id1 = ids.GenerateTestID()
    17  	id2 = ids.GenerateTestID()
    18  	id3 = ids.GenerateTestID()
    19  	id4 = ids.GenerateTestID()
    20  )
    21  
    22  func TestAdd(t *testing.T) {
    23  	tests := map[string]struct {
    24  		initial  Tree
    25  		blkID    ids.ID
    26  		parentID ids.ID
    27  		expected Tree
    28  	}{
    29  		"add to empty tree": {
    30  			initial: &tree{
    31  				childToParent:    make(map[ids.ID]ids.ID),
    32  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
    33  			},
    34  			blkID:    id1,
    35  			parentID: id2,
    36  			expected: &tree{
    37  				childToParent: map[ids.ID]ids.ID{
    38  					id1: id2,
    39  				},
    40  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
    41  					id2: set.Of(id1),
    42  				},
    43  			},
    44  		},
    45  		"add new parent": {
    46  			initial: &tree{
    47  				childToParent: map[ids.ID]ids.ID{
    48  					id1: id2,
    49  				},
    50  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
    51  					id2: set.Of(id1),
    52  				},
    53  			},
    54  			blkID:    id3,
    55  			parentID: id4,
    56  			expected: &tree{
    57  				childToParent: map[ids.ID]ids.ID{
    58  					id1: id2,
    59  					id3: id4,
    60  				},
    61  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
    62  					id2: set.Of(id1),
    63  					id4: set.Of(id3),
    64  				},
    65  			},
    66  		},
    67  		"add new block to existing parent": {
    68  			initial: &tree{
    69  				childToParent: map[ids.ID]ids.ID{
    70  					id1: id2,
    71  				},
    72  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
    73  					id2: set.Of(id1),
    74  				},
    75  			},
    76  			blkID:    id3,
    77  			parentID: id2,
    78  			expected: &tree{
    79  				childToParent: map[ids.ID]ids.ID{
    80  					id1: id2,
    81  					id3: id2,
    82  				},
    83  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
    84  					id2: set.Of(id1, id3),
    85  				},
    86  			},
    87  		},
    88  	}
    89  	for name, test := range tests {
    90  		t.Run(name, func(t *testing.T) {
    91  			at := test.initial
    92  			at.Add(test.blkID, test.parentID)
    93  			require.Equal(t, test.expected, at)
    94  		})
    95  	}
    96  }
    97  
    98  func TestRemove(t *testing.T) {
    99  	tests := map[string]struct {
   100  		initial  Tree
   101  		blkID    ids.ID
   102  		expected Tree
   103  	}{
   104  		"remove from empty tree": {
   105  			initial: &tree{
   106  				childToParent:    make(map[ids.ID]ids.ID),
   107  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   108  			},
   109  			blkID: id1,
   110  			expected: &tree{
   111  				childToParent:    make(map[ids.ID]ids.ID),
   112  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   113  			},
   114  		},
   115  		"remove block and parent from tree": {
   116  			initial: &tree{
   117  				childToParent: map[ids.ID]ids.ID{
   118  					id1: id2,
   119  				},
   120  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   121  					id2: set.Of(id1),
   122  				},
   123  			},
   124  			blkID: id1,
   125  			expected: &tree{
   126  				childToParent:    make(map[ids.ID]ids.ID),
   127  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   128  			},
   129  		},
   130  		"remove block and not parent from tree": {
   131  			initial: &tree{
   132  				childToParent: map[ids.ID]ids.ID{
   133  					id1: id2,
   134  					id3: id2,
   135  				},
   136  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   137  					id2: set.Of(id1, id3),
   138  				},
   139  			},
   140  			blkID: id1,
   141  			expected: &tree{
   142  				childToParent: map[ids.ID]ids.ID{
   143  					id3: id2,
   144  				},
   145  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   146  					id2: set.Of(id3),
   147  				},
   148  			},
   149  		},
   150  		"remove untracked block from tree": {
   151  			initial: &tree{
   152  				childToParent: map[ids.ID]ids.ID{
   153  					id1: id2,
   154  				},
   155  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   156  					id2: set.Of(id1),
   157  				},
   158  			},
   159  			blkID: id2,
   160  			expected: &tree{
   161  				childToParent: map[ids.ID]ids.ID{
   162  					id1: id2,
   163  				},
   164  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   165  					id2: set.Of(id1),
   166  				},
   167  			},
   168  		},
   169  	}
   170  	for name, test := range tests {
   171  		t.Run(name, func(t *testing.T) {
   172  			at := test.initial
   173  			at.Remove(test.blkID)
   174  			require.Equal(t, test.expected, at)
   175  		})
   176  	}
   177  }
   178  
   179  func TestRemoveDescendants(t *testing.T) {
   180  	tests := map[string]struct {
   181  		initial  Tree
   182  		blkID    ids.ID
   183  		expected Tree
   184  	}{
   185  		"remove from empty tree": {
   186  			initial: &tree{
   187  				childToParent:    make(map[ids.ID]ids.ID),
   188  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   189  			},
   190  			blkID: id1,
   191  			expected: &tree{
   192  				childToParent:    make(map[ids.ID]ids.ID),
   193  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   194  			},
   195  		},
   196  		"remove block and parent from tree": {
   197  			initial: &tree{
   198  				childToParent: map[ids.ID]ids.ID{
   199  					id1: id2,
   200  				},
   201  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   202  					id2: set.Of(id1),
   203  				},
   204  			},
   205  			blkID: id1,
   206  			expected: &tree{
   207  				childToParent:    make(map[ids.ID]ids.ID),
   208  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   209  			},
   210  		},
   211  		"remove block and not parent from tree": {
   212  			initial: &tree{
   213  				childToParent: map[ids.ID]ids.ID{
   214  					id1: id2,
   215  					id3: id2,
   216  				},
   217  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   218  					id2: set.Of(id1, id3),
   219  				},
   220  			},
   221  			blkID: id1,
   222  			expected: &tree{
   223  				childToParent: map[ids.ID]ids.ID{
   224  					id3: id2,
   225  				},
   226  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   227  					id2: set.Of(id3),
   228  				},
   229  			},
   230  		},
   231  		"remove untracked block from tree": {
   232  			initial: &tree{
   233  				childToParent: map[ids.ID]ids.ID{
   234  					id1: id2,
   235  				},
   236  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   237  					id2: set.Of(id1),
   238  				},
   239  			},
   240  			blkID: id3,
   241  			expected: &tree{
   242  				childToParent: map[ids.ID]ids.ID{
   243  					id1: id2,
   244  				},
   245  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   246  					id2: set.Of(id1),
   247  				},
   248  			},
   249  		},
   250  		"remove children from tree": {
   251  			initial: &tree{
   252  				childToParent: map[ids.ID]ids.ID{
   253  					id1: id2,
   254  					id3: id2,
   255  				},
   256  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   257  					id2: set.Of(id1, id3),
   258  				},
   259  			},
   260  			blkID: id2,
   261  			expected: &tree{
   262  				childToParent:    make(map[ids.ID]ids.ID),
   263  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   264  			},
   265  		},
   266  		"remove grand child from tree": {
   267  			initial: &tree{
   268  				childToParent: map[ids.ID]ids.ID{
   269  					id1: id2,
   270  					id2: id3,
   271  				},
   272  				parentToChildren: map[ids.ID]set.Set[ids.ID]{
   273  					id2: set.Of(id1),
   274  					id3: set.Of(id2),
   275  				},
   276  			},
   277  			blkID: id3,
   278  			expected: &tree{
   279  				childToParent:    make(map[ids.ID]ids.ID),
   280  				parentToChildren: make(map[ids.ID]set.Set[ids.ID]),
   281  			},
   282  		},
   283  	}
   284  	for name, test := range tests {
   285  		t.Run(name, func(t *testing.T) {
   286  			at := test.initial
   287  			at.RemoveDescendants(test.blkID)
   288  			require.Equal(t, test.expected, at)
   289  		})
   290  	}
   291  }
   292  
   293  func TestHas(t *testing.T) {
   294  	require := require.New(t)
   295  
   296  	at := NewTree()
   297  	require.False(at.Has(id1))
   298  	require.False(at.Has(id2))
   299  	require.False(at.Has(id3))
   300  
   301  	at.Add(id1, id2)
   302  	require.True(at.Has(id1))
   303  	require.False(at.Has(id2))
   304  	require.False(at.Has(id3))
   305  
   306  	at.Add(id2, id3)
   307  	require.True(at.Has(id1))
   308  	require.True(at.Has(id2))
   309  	require.False(at.Has(id3))
   310  }
   311  
   312  func TestGetAncestor(t *testing.T) {
   313  	require := require.New(t)
   314  
   315  	at := NewTree()
   316  	require.Equal(id1, at.GetAncestor(id1))
   317  	require.Equal(id2, at.GetAncestor(id2))
   318  	require.Equal(id3, at.GetAncestor(id3))
   319  	require.Equal(id4, at.GetAncestor(id4))
   320  
   321  	at.Add(id1, id2)
   322  	require.Equal(id2, at.GetAncestor(id1))
   323  	require.Equal(id2, at.GetAncestor(id2))
   324  	require.Equal(id3, at.GetAncestor(id3))
   325  	require.Equal(id4, at.GetAncestor(id4))
   326  
   327  	at.Add(id2, id3)
   328  	require.Equal(id3, at.GetAncestor(id1))
   329  	require.Equal(id3, at.GetAncestor(id2))
   330  	require.Equal(id3, at.GetAncestor(id3))
   331  	require.Equal(id4, at.GetAncestor(id4))
   332  
   333  	at.Add(id4, id3)
   334  	require.Equal(id3, at.GetAncestor(id1))
   335  	require.Equal(id3, at.GetAncestor(id2))
   336  	require.Equal(id3, at.GetAncestor(id3))
   337  	require.Equal(id3, at.GetAncestor(id4))
   338  }
   339  
   340  func TestLen(t *testing.T) {
   341  	require := require.New(t)
   342  
   343  	at := NewTree()
   344  	require.Zero(at.Len())
   345  
   346  	at.Add(id1, id2)
   347  	require.Equal(1, at.Len())
   348  
   349  	at.Add(id2, id3)
   350  	require.Equal(2, at.Len())
   351  
   352  	at.Add(id4, id3)
   353  	require.Equal(3, at.Len())
   354  }