partition_movement_test: fix offset key error

Occassionally we would hit errors like the following:

RunnerClient: rptest.tests.partition_movement_test.PartitionMovementTest.test_bootstrapping_after_move.num_to_upgrade=2: Summary: KeyError(1)
Traceback (most recent call last):
  File "/usr/local/lib/python3.9/dist-packages/ducktape/tests/runner_client.py", line 135, in run
    data = self.run_test()
  File "/usr/local/lib/python3.9/dist-packages/ducktape/tests/runner_client.py", line 227, in run_test
    return self.test_context.function(self.test)
  File "/usr/local/lib/python3.9/dist-packages/ducktape/mark/_mark.py", line 476, in wrapper
    return functools.partial(f, *args, **kwargs)(*w_args, **w_kwargs)
  File "/root/tests/rptest/services/cluster.py", line 35, in wrapped
    r = f(self, *args, **kwargs)
  File "/root/tests/rptest/tests/partition_movement_test.py", line 410, in test_bootstrapping_after_move
    wait_until(offsets_are_recovered, 30, 2)
  File "/usr/local/lib/python3.9/dist-packages/ducktape/utils/util.py", line 53, in wait_until
    raise e
  File "/usr/local/lib/python3.9/dist-packages/ducktape/utils/util.py", line 44, in wait_until
    if condition():
  File "/root/tests/rptest/tests/partition_movement_test.py", line 405, in offsets_are_recovered
    return all([
  File "/root/tests/rptest/tests/partition_movement_test.py", line 406, in <listcomp>
    offset_map[p.id] == p.high_watermark
KeyError: 1

It's possible that the command run by `describe_topic()` results in
lines that don't match the partition line regex. The regex expects
numeric values for each field, which isn't always the case e.g. if
there's a leadership change.

This commit adjusts the test to ensure there are three partitions in the
returned output before proceeding.

tests/rptest/tests/partition_movement_test.py

@@ -374,7 +374,19 @@ def test_bootstrapping_after_move(self, num_to_upgrade):
 
         # snapshot offsets
         rpk = RpkTool(self.redpanda)
-        partitions = rpk.describe_topic(spec.name)
+
+        def has_offsets_for_all_partitions(out):
+            # NOTE: partitions may not be returned if their fields can't be
+            # populated, e.g. during leadership changes.
+            partitions = list(rpk.describe_topic(spec.name))
+            if len(partitions) == 3:
+                out += partitions
+                return True
+            return False
+
+        partitions = []
+        wait_until(lambda: has_offsets_for_all_partitions(partitions), 30, 2)
+
         offset_map = {}
         for p in partitions:
             offset_map[p.id] = p.high_watermark
@@ -383,10 +395,11 @@ def test_bootstrapping_after_move(self, num_to_upgrade):
         self.redpanda.restart_nodes(self.redpanda.nodes)
 
         def offsets_are_recovered():
-
+            partitions_after = list(rpk.describe_topic(spec.name))
+            if len(partitions_after) != 3:
+                return False
             return all([
-                offset_map[p.id] == p.high_watermark
-                for p in rpk.describe_topic(spec.name)
+                offset_map[p.id] == p.high_watermark for p in partitions_after
             ])
 
         wait_until(offsets_are_recovered, 30, 2)