wtclient: divide addr_iter tests into subtests

Split out the existing tests into logical subtests.

watchtower/wtclient/addr_iterator_test.go

@@ -10,184 +10,210 @@ import (
 
 // TestAddrIterator tests the behaviour of the addressIterator.
 func TestAddrIterator(t *testing.T) {
-	// Assert that an iterator can't be initialised with an empty address
+	t.Parallel()
-	// list.
-	_, err := newAddressIterator()
-	require.ErrorContains(t, err, "must have at least one address")
 
+	// Prep some addresses to use in the subtests.
 	addr1, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8000")
 	require.NoError(t, err)
 
-	// Initialise the iterator with addr1.
-	iter, err := newAddressIterator(addr1)
-	require.NoError(t, err)
-
-	// Attempting to remove addr1 should fail now since it is the only
-	// address in the iterator.
-	iter.Add(addr1)
-	err = iter.Remove(addr1)
-	require.ErrorIs(t, err, wtdb.ErrLastTowerAddr)
-
-	// Adding a duplicate of addr1 and then calling Remove should still
-	// return an error.
-	err = iter.Remove(addr1)
-	require.ErrorIs(t, err, wtdb.ErrLastTowerAddr)
-
 	addr2, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8001")
 	require.NoError(t, err)
 
-	// Add addr2 to the iterator.
-	iter.Add(addr2)
-
-	// Check that peek returns addr1.
-	a1 := iter.Peek()
-	require.NoError(t, err)
-	require.Equal(t, addr1, a1)
-
-	// Calling peek multiple times should return the same result.
-	a1 = iter.Peek()
-	require.Equal(t, addr1, a1)
-
-	// Calling Next should now return addr2.
-	a2, err := iter.Next()
-	require.NoError(t, err)
-	require.Equal(t, addr2, a2)
-
-	// Assert that Peek now returns addr2.
-	a2 = iter.Peek()
-	require.NoError(t, err)
-	require.Equal(t, addr2, a2)
-
-	// Calling Next should result in reaching the end of th list.
-	_, err = iter.Next()
-	require.ErrorIs(t, err, ErrAddressesExhausted)
-
-	// Calling Peek now should reset the queue and return addr1.
-	a1 = iter.Peek()
-	require.Equal(t, addr1, a1)
-
-	// Wind the list to the end again so that we can test the Reset func.
-	_, err = iter.Next()
-	require.NoError(t, err)
-
-	_, err = iter.Next()
-	require.ErrorIs(t, err, ErrAddressesExhausted)
-
-	iter.Reset()
-
-	// Now Next should return addr 2.
-	a2, err = iter.Next()
-	require.NoError(t, err)
-	require.Equal(t, addr2, a2)
-
 	addr3, err := net.ResolveTCPAddr("tcp", "1.2.3.4:8002")
 	require.NoError(t, err)
 
-	// Add addr3 now to ensure that the iteration works even if we are
+	t.Run("always at least one address", func(t *testing.T) {
-	// midway through the queue.
+		t.Parallel()
-	iter.Add(addr3)
 
-	// Now Next should return addr 3.
+		// Assert that an iterator can't be initialised with an empty
-	a3, err := iter.Next()
+		// address list.
-	require.NoError(t, err)
+		_, err := newAddressIterator()
-	require.Equal(t, addr3, a3)
+		require.ErrorContains(t, err, "must have at least one address")
 
-	// Quickly test that GetAll correctly returns a copy of all the
+		// Initialise the iterator with addr1.
-	// addresses in the iterator.
+		iter, err := newAddressIterator(addr1)
-	addrList := iter.GetAll()
+		require.NoError(t, err)
-	require.ElementsMatch(t, addrList, []net.Addr{addr1, addr2, addr3})
 
-	// Also check that an iterator constructed via the Copy method, also
+		// Attempting to remove addr1 should fail now since it is the
-	// contains all the expected addresses.
+		// only address in the iterator.
-	newIterAddrs := iter.Copy().GetAll()
+		iter.Add(addr1)
-	require.ElementsMatch(t, newIterAddrs, []net.Addr{addr1, addr2, addr3})
+		err = iter.Remove(addr1)
+		require.ErrorIs(t, err, wtdb.ErrLastTowerAddr)
 
-	// Let's now remove addr3.
+		// Adding a duplicate of addr1 and then calling Remove should
-	err = iter.Remove(addr3)
+		// still return an error.
-	require.NoError(t, err)
+		err = iter.Remove(addr1)
+		require.ErrorIs(t, err, wtdb.ErrLastTowerAddr)
+	})
 
-	// Since addr3 is gone, Peek should return addr1.
+	t.Run("address iteration", func(t *testing.T) {
-	a1 = iter.Peek()
+		t.Parallel()
-	require.Equal(t, addr1, a1)
 
-	// Lastly, we will test the "locking" of addresses.
+		// Initialise the iterator with addr1.
+		iter, err := newAddressIterator(addr1)
+		require.NoError(t, err)
 
-	// First we test the locking of an address via the PeekAndLock function.
+		// Add addr2 to the iterator.
-	a1 = iter.PeekAndLock()
+		iter.Add(addr2)
-	require.Equal(t, addr1, a1)
-	require.True(t, iter.HasLocked())
 
-	// Assert that we can't remove addr1 if there is a lock on it.
+		// Check that peek returns addr1.
-	err = iter.Remove(addr1)
+		a1 := iter.Peek()
-	require.ErrorIs(t, err, ErrAddrInUse)
+		require.NoError(t, err)
+		require.Equal(t, addr1, a1)
 
-	// Now release the lock on addr1.
+		// Calling peek multiple times should return the same result.
-	iter.ReleaseLock(addr1)
+		a1 = iter.Peek()
-	require.False(t, iter.HasLocked())
+		require.Equal(t, addr1, a1)
 
-	// Since the lock has been released, we should now be able to remove
+		// Calling Next should now return addr2.
-	// addr1.
+		a2, err := iter.Next()
-	err = iter.Remove(addr1)
+		require.NoError(t, err)
-	require.NoError(t, err)
+		require.Equal(t, addr2, a2)
 
-	// Now we test the locking of an address via the NextAndLock function.
+		// Assert that Peek now returns addr2.
-	// To do this, we first re-add addr3.
+		a2 = iter.Peek()
-	iter.Add(addr3)
+		require.NoError(t, err)
+		require.Equal(t, addr2, a2)
 
-	a2, err = iter.NextAndLock()
+		// Calling Next should result in reaching the end of th list.
-	require.NoError(t, err)
+		_, err = iter.Next()
-	require.Equal(t, addr2, a2)
+		require.ErrorIs(t, err, ErrAddressesExhausted)
-	require.True(t, iter.HasLocked())
 
-	// Assert that we can't remove addr2 if there is a lock on it.
+		// Calling Peek now should reset the queue and return addr1.
-	err = iter.Remove(addr2)
+		a1 = iter.Peek()
-	require.ErrorIs(t, err, ErrAddrInUse)
+		require.Equal(t, addr1, a1)
 
-	// Now release the lock on addr2.
+		// Wind the list to the end again so that we can test the Reset
-	iter.ReleaseLock(addr2)
+		// func.
-	require.False(t, iter.HasLocked())
+		_, err = iter.Next()
+		require.NoError(t, err)
 
-	// Since the lock has been released, we should now be able to remove
+		_, err = iter.Next()
-	// addr1.
+		require.ErrorIs(t, err, ErrAddressesExhausted)
-	err = iter.Remove(addr2)
-	require.NoError(t, err)
 
-	// Only addr3 should still be left in the iterator.
+		iter.Reset()
-	addrList = iter.GetAll()
-	require.Len(t, addrList, 1)
-	require.Contains(t, addrList, addr3)
 
-	// Ensure that HasLocked acts correctly in the case where more than one
+		// Now Next should return addr 2.
-	// address is being locked and unlock as well as the case where the same
+		a2, err = iter.Next()
-	// address is locked more than once.
+		require.NoError(t, err)
+		require.Equal(t, addr2, a2)
 
-	require.False(t, iter.HasLocked())
+		// Add addr3 now to ensure that the iteration works even if we
+		// are midway through the queue.
+		iter.Add(addr3)
 
-	a3 = iter.PeekAndLock()
+		// Now Next should return addr 3.
-	require.Equal(t, addr3, a3)
+		a3, err := iter.Next()
-	require.True(t, iter.HasLocked())
+		require.NoError(t, err)
+		require.Equal(t, addr3, a3)
 
-	a3 = iter.PeekAndLock()
+		// Quickly test that GetAll correctly returns a copy of all the
-	require.Equal(t, addr3, a3)
+		// addresses in the iterator.
-	require.True(t, iter.HasLocked())
+		addrList := iter.GetAll()
+		require.ElementsMatch(
+			t, addrList, []net.Addr{addr1, addr2, addr3},
+		)
 
-	iter.Add(addr2)
+		// Also check that an iterator constructed via the Copy method,
-	a2, err = iter.NextAndLock()
+		// also contains all the expected addresses.
-	require.NoError(t, err)
+		newIterAddrs := iter.Copy().GetAll()
-	require.Equal(t, addr2, a2)
+		require.ElementsMatch(
-	require.True(t, iter.HasLocked())
+			t, newIterAddrs, []net.Addr{addr1, addr2, addr3},
+		)
 
-	// Now release addr2 and asset that HasLock is still true.
+		// Let's now remove addr3.
-	iter.ReleaseLock(addr2)
+		err = iter.Remove(addr3)
-	require.True(t, iter.HasLocked())
+		require.NoError(t, err)
 
-	// Releasing one of the locks on addr3 now should still result in
+		// Since addr3 is gone, Peek should return addr1.
-	// HasLocked returning true.
+		a1 = iter.Peek()
-	iter.ReleaseLock(addr3)
+		require.Equal(t, addr1, a1)
-	require.True(t, iter.HasLocked())
+	})
 
-	// Releasing it again should now result in should still result in
+	t.Run("address locking", func(t *testing.T) {
-	// HasLocked returning false.
+		t.Parallel()
-	iter.ReleaseLock(addr3)
+
-	require.False(t, iter.HasLocked())
+		// Initialise the iterator with addr1 and addr2.
+		iter, err := newAddressIterator(addr1, addr2)
+		require.NoError(t, err)
+
+		// First we test the locking of an address via the PeekAndLock
+		// function.
+		a1 := iter.PeekAndLock()
+		require.Equal(t, addr1, a1)
+		require.True(t, iter.HasLocked())
+
+		// Assert that we can't remove addr1 if there is a lock on it.
+		err = iter.Remove(addr1)
+		require.ErrorIs(t, err, ErrAddrInUse)
+
+		// Now release the lock on addr1.
+		iter.ReleaseLock(addr1)
+		require.False(t, iter.HasLocked())
+
+		// Since the lock has been released, we should now be able to
+		// remove addr1.
+		err = iter.Remove(addr1)
+		require.NoError(t, err)
+
+		// Now we test the locking of an address via the NextAndLock
+		// function. To do this, we first re-add addr3.
+		iter.Add(addr3)
+
+		a2, err := iter.NextAndLock()
+		require.NoError(t, err)
+		require.Equal(t, addr2, a2)
+		require.True(t, iter.HasLocked())
+
+		// Assert that we can't remove addr2 if there is a lock on it.
+		err = iter.Remove(addr2)
+		require.ErrorIs(t, err, ErrAddrInUse)
+
+		// Now release the lock on addr2.
+		iter.ReleaseLock(addr2)
+		require.False(t, iter.HasLocked())
+
+		// Since the lock has been released, we should now be able to
+		// remove addr1.
+		err = iter.Remove(addr2)
+		require.NoError(t, err)
+
+		// Only addr3 should still be left in the iterator.
+		addrList := iter.GetAll()
+		require.Len(t, addrList, 1)
+		require.Contains(t, addrList, addr3)
+
+		// Ensure that HasLocked acts correctly in the case where more
+		// than one address is being locked and unlock as well as the
+		// case where the same address is locked more than once.
+		require.False(t, iter.HasLocked())
+
+		a3 := iter.PeekAndLock()
+		require.Equal(t, addr3, a3)
+		require.True(t, iter.HasLocked())
+
+		a3 = iter.PeekAndLock()
+		require.Equal(t, addr3, a3)
+		require.True(t, iter.HasLocked())
+
+		iter.Add(addr2)
+		a2, err = iter.NextAndLock()
+		require.NoError(t, err)
+		require.Equal(t, addr2, a2)
+		require.True(t, iter.HasLocked())
+
+		// Now release addr2 and assert that HasLock is still true.
+		iter.ReleaseLock(addr2)
+		require.True(t, iter.HasLocked())
+
+		// Releasing one of the locks on addr3 now should still result
+		// in HasLocked returning true.
+		iter.ReleaseLock(addr3)
+		require.True(t, iter.HasLocked())
+
+		// Releasing it again should now result in should still result
+		// in HasLocked returning false.
+		iter.ReleaseLock(addr3)
+		require.False(t, iter.HasLocked())
+	})
 }