diff --git a/netmap/subnet.go b/netmap/subnet.go index 09b7a57..b61a504 100644 --- a/netmap/subnet.go +++ b/netmap/subnet.go @@ -11,6 +11,15 @@ import ( // EnterSubnet writes to NodeInfo the intention to enter the subnet. Must not be called on nil. // Zero NodeInfo belongs to zero subnet. func (i *NodeInfo) EnterSubnet(id subnetid.ID) { + i.changeSubnet(id, true) +} + +// ExitSubnet writes to NodeInfo the intention to exit subnet. Must not be called on nil. +func (i *NodeInfo) ExitSubnet(id subnetid.ID) { + i.changeSubnet(id, false) +} + +func (i *NodeInfo) changeSubnet(id subnetid.ID, isMember bool) { var ( idv2 refs.SubnetID info netmap.NodeSubnetInfo @@ -19,7 +28,7 @@ func (i *NodeInfo) EnterSubnet(id subnetid.ID) { id.WriteToV2(&idv2) info.SetID(&idv2) - info.SetEntryFlag(true) + info.SetEntryFlag(isMember) netmap.WriteSubnetInfo((*netmap.NodeInfo)(i), info) } diff --git a/netmap/subnet_test.go b/netmap/subnet_test.go index cb1c0c4..0dd3600 100644 --- a/netmap/subnet_test.go +++ b/netmap/subnet_test.go @@ -86,6 +86,30 @@ func TestNodeInfoSubnets(t *testing.T) { }) } +func TestEnterSubnet(t *testing.T) { + var ( + id subnetid.ID + node netmap.NodeInfo + ) + + require.True(t, netmap.BelongsToSubnet(&node, id)) + + node.EnterSubnet(id) + require.True(t, netmap.BelongsToSubnet(&node, id)) + + node.ExitSubnet(id) + require.False(t, netmap.BelongsToSubnet(&node, id)) + + id.SetNumber(10) + node.EnterSubnet(id) + require.True(t, netmap.BelongsToSubnet(&node, id)) + require.False(t, netmap.BelongsToSubnet(&node, subnetid.ID{})) + + node.ExitSubnet(id) + require.False(t, netmap.BelongsToSubnet(&node, id)) + require.False(t, netmap.BelongsToSubnet(&node, subnetid.ID{})) +} + func TestBelongsToSubnet(t *testing.T) { var id, idMiss, idZero subnetid.ID