More iterator tests

query/iter_test.go

@@ -200,3 +200,61 @@ func TestClosestNodesIterTimedOutNodeMakesCapacity(t *testing.T) {
 	state = iter.Advance(ctx, nil)
 	require.IsType(t, &StateNodeIterWaitingWithCapacity{}, state)
 }
+
+func TestClosestNodesIterSuccessfulContactMakesCapacity(t *testing.T) {
+	ctx := context.Background()
+
+	target := key.Key8(0b00000001)
+	a := kadtest.NewID(key.Key8(0b00000100)) // 4
+	b := kadtest.NewID(key.Key8(0b00001000)) // 8
+	c := kadtest.NewID(key.Key8(0b00010000)) // 16
+	d := kadtest.NewID(key.Key8(0b00100000)) // 32
+
+	// ensure the order of the known nodes
+	require.True(t, target.Xor(a.Key()).Compare(target.Xor(b.Key())) == -1)
+	require.True(t, target.Xor(b.Key()).Compare(target.Xor(c.Key())) == -1)
+	require.True(t, target.Xor(c.Key()).Compare(target.Xor(d.Key())) == -1)
+
+	// knownNodes are in "random" order
+	knownNodes := []kad.NodeID[key.Key8]{b, c, a, d}
+
+	clk := clock.NewMock()
+
+	cfg := DefaultClosestNodesIterConfig()
+	cfg.Clock = clk
+	cfg.Concurrency = len(knownNodes) - 1 // one less than the number of initial nodes
+
+	iter := NewClosestNodesIter(target, knownNodes, cfg)
+
+	// first thing the new iterator should do is contact the nearest node
+	state := iter.Advance(ctx, nil)
+	require.IsType(t, &StateNodeIterWaitingContact[key.Key8]{}, state)
+	st := state.(*StateNodeIterWaitingContact[key.Key8])
+	require.Equal(t, a, st.NodeID)
+
+	// while the iter has capacity the iter should contact the next nearest node
+	state = iter.Advance(ctx, nil)
+	require.IsType(t, &StateNodeIterWaitingContact[key.Key8]{}, state)
+	st = state.(*StateNodeIterWaitingContact[key.Key8])
+	require.Equal(t, b, st.NodeID)
+
+	// while the iter has capacity the iter should contact the second nearest node
+	state = iter.Advance(ctx, nil)
+	require.IsType(t, &StateNodeIterWaitingContact[key.Key8]{}, state)
+	st = state.(*StateNodeIterWaitingContact[key.Key8])
+	require.Equal(t, c, st.NodeID)
+
+	// the iter should be at capacity
+	state = iter.Advance(ctx, nil)
+	require.IsType(t, &StateNodeIterWaitingAtCapacity{}, state)
+
+	// notify iterator that first node was contacted successfully, now node d can be contacted
+	state = iter.Advance(ctx, &EventNodeIterNodeContacted[key.Key8]{NodeID: a})
+	require.IsType(t, &StateNodeIterWaitingContact[key.Key8]{}, state)
+	st = state.(*StateNodeIterWaitingContact[key.Key8])
+	require.Equal(t, d, st.NodeID)
+
+	// the iter should be at capacity again
+	state = iter.Advance(ctx, nil)
+	require.IsType(t, &StateNodeIterWaitingAtCapacity{}, state)
+}