usb: Fix test failure with multiple partitions

When test_usb_load finds multiple partitions of the same type then
it will cause a test failure. The call to write the test file will
write a different test file to each partition but only return the
name and size of the last one written. So the test then fails to
load the test file from the first partition as it uses the name of
a file on a different partition.

Refactor the code so that only one test file is written at a time
and is written to only the partition being tested at that time. This
allows the correct file name to always be available to the code that
runs the load command. This reduces the number of files written and
also the number of calls to crc32 needed.

Fixes: 1c5b6edad381 ("test/py: usb: Add tests for USB device")
Signed-off-by: Andrew Goodbody <andrew.goodbody@linaro.org>
Reviewed-by: Love Kumar <love.kumar@amd.com>
diff --git a/test/py/tests/test_usb.py b/test/py/tests/test_usb.py
index fb3d20f..e7966ca 100644
--- a/test/py/tests/test_usb.py
+++ b/test/py/tests/test_usb.py
@@ -288,6 +288,47 @@
     if not part_detect:
         pytest.skip('No %s partition detected' % fs.upper())
 
+def usb_fatload_fatwrite(u_boot_console, fs, x, part):
+    addr = u_boot_utils.find_ram_base(u_boot_console)
+    size = random.randint(4, 1 * 1024 * 1024)
+    output = u_boot_console.run_command('crc32 %x %x' % (addr, size))
+    m = re.search('==> (.+?)', output)
+    if not m:
+        pytest.fail('CRC32 failed')
+    expected_crc32 = m.group(1)
+
+    file = '%s_%d' % ('uboot_test', size)
+    output = u_boot_console.run_command(
+        '%swrite usb %d:%s %x %s %x' % (fs, x, part, addr, file, size)
+    )
+    assert 'Unable to write' not in output
+    assert 'Error' not in output
+    assert 'overflow' not in output
+    expected_text = '%d bytes written' % size
+    assert expected_text in output
+
+    alignment = int(
+        u_boot_console.config.buildconfig.get(
+            'config_sys_cacheline_size', 128
+        )
+    )
+    offset = random.randrange(alignment, 1024, alignment)
+    output = u_boot_console.run_command(
+        '%sload usb %d:%s %x %s' % (fs, x, part, addr + offset, file)
+    )
+    assert 'Invalid FAT entry' not in output
+    assert 'Unable to read file' not in output
+    assert 'Misaligned buffer address' not in output
+    expected_text = '%d bytes read' % size
+    assert expected_text in output
+
+    output = u_boot_console.run_command(
+        'crc32 %x $filesize' % (addr + offset)
+    )
+    assert expected_crc32 in output
+
+    return file, size, expected_crc32
+
 @pytest.mark.buildconfigspec('cmd_usb')
 @pytest.mark.buildconfigspec('cmd_fat')
 @pytest.mark.buildconfigspec('cmd_memory')
@@ -309,49 +350,11 @@
 
             for part in partitions:
                 part_detect = 1
-                addr = u_boot_utils.find_ram_base(u_boot_console)
-                size = random.randint(4, 1 * 1024 * 1024)
-                output = u_boot_console.run_command('crc32 %x %x' % (addr, size))
-                m = re.search('==> (.+?)', output)
-                if not m:
-                    pytest.fail('CRC32 failed')
-                expected_crc32 = m.group(1)
-
-                file = '%s_%d' % ('uboot_test', size)
-                output = u_boot_console.run_command(
-                    '%swrite usb %d:%s %x %s %x' % (fs, x, part, addr, file, size)
-                )
-                assert 'Unable to write' not in output
-                assert 'Error' not in output
-                assert 'overflow' not in output
-                expected_text = '%d bytes written' % size
-                assert expected_text in output
-
-                alignment = int(
-                    u_boot_console.config.buildconfig.get(
-                        'config_sys_cacheline_size', 128
-                    )
-                )
-                offset = random.randrange(alignment, 1024, alignment)
-                output = u_boot_console.run_command(
-                    '%sload usb %d:%s %x %s' % (fs, x, part, addr + offset, file)
-                )
-                assert 'Invalid FAT entry' not in output
-                assert 'Unable to read file' not in output
-                assert 'Misaligned buffer address' not in output
-                expected_text = '%d bytes read' % size
-                assert expected_text in output
-
-                output = u_boot_console.run_command(
-                    'crc32 %x $filesize' % (addr + offset)
-                )
-                assert expected_crc32 in output
+                usb_fatload_fatwrite(u_boot_console, fs, x, part)
 
     if not part_detect:
         pytest.skip('No %s partition detected' % fs.upper())
 
-    return file, size
-
 @pytest.mark.buildconfigspec('cmd_usb')
 @pytest.mark.buildconfigspec('cmd_ext4')
 def test_usb_ext4ls(u_boot_console):
@@ -380,6 +383,39 @@
     if not part_detect:
         pytest.skip('No %s partition detected' % fs.upper())
 
+def usb_ext4load_ext4write(u_boot_console, fs, x, part):
+    addr = u_boot_utils.find_ram_base(u_boot_console)
+    size = random.randint(4, 1 * 1024 * 1024)
+    output = u_boot_console.run_command('crc32 %x %x' % (addr, size))
+    m = re.search('==> (.+?)', output)
+    if not m:
+        pytest.fail('CRC32 failed')
+    expected_crc32 = m.group(1)
+    file = '%s_%d' % ('uboot_test', size)
+
+    output = u_boot_console.run_command(
+        '%swrite usb %d:%s %x /%s %x' % (fs, x, part, addr, file, size)
+    )
+    assert 'Unable to write' not in output
+    assert 'Error' not in output
+    assert 'overflow' not in output
+    expected_text = '%d bytes written' % size
+    assert expected_text in output
+
+    offset = random.randrange(128, 1024, 128)
+    output = u_boot_console.run_command(
+        '%sload usb %d:%s %x /%s' % (fs, x, part, addr + offset, file)
+    )
+    expected_text = '%d bytes read' % size
+    assert expected_text in output
+
+    output = u_boot_console.run_command(
+        'crc32 %x $filesize' % (addr + offset)
+    )
+    assert expected_crc32 in output
+
+    return file, size, expected_crc32
+
 @pytest.mark.buildconfigspec('cmd_usb')
 @pytest.mark.buildconfigspec('cmd_ext4')
 @pytest.mark.buildconfigspec('ext4_write')
@@ -402,41 +438,11 @@
 
             for part in partitions:
                 part_detect = 1
-                addr = u_boot_utils.find_ram_base(u_boot_console)
-                size = random.randint(4, 1 * 1024 * 1024)
-                output = u_boot_console.run_command('crc32 %x %x' % (addr, size))
-                m = re.search('==> (.+?)', output)
-                if not m:
-                    pytest.fail('CRC32 failed')
-                expected_crc32 = m.group(1)
-                file = '%s_%d' % ('uboot_test', size)
-
-                output = u_boot_console.run_command(
-                    '%swrite usb %d:%s %x /%s %x' % (fs, x, part, addr, file, size)
-                )
-                assert 'Unable to write' not in output
-                assert 'Error' not in output
-                assert 'overflow' not in output
-                expected_text = '%d bytes written' % size
-                assert expected_text in output
-
-                offset = random.randrange(128, 1024, 128)
-                output = u_boot_console.run_command(
-                    '%sload usb %d:%s %x /%s' % (fs, x, part, addr + offset, file)
-                )
-                expected_text = '%d bytes read' % size
-                assert expected_text in output
-
-                output = u_boot_console.run_command(
-                    'crc32 %x $filesize' % (addr + offset)
-                )
-                assert expected_crc32 in output
+                usb_ext4load_ext4write(u_boot_console, fs, x, part)
 
     if not part_detect:
         pytest.skip('No %s partition detected' % fs.upper())
 
-    return file, size
-
 @pytest.mark.buildconfigspec('cmd_usb')
 @pytest.mark.buildconfigspec('cmd_ext2')
 def test_usb_ext2ls(u_boot_console):
@@ -473,7 +479,6 @@
 @pytest.mark.buildconfigspec('cmd_memory')
 def test_usb_ext2load(u_boot_console):
     devices, controllers, storage_device = test_usb_part(u_boot_console)
-    file, size = test_usb_ext4load_ext4write(u_boot_console)
 
     if not devices:
         pytest.skip('No devices detected')
@@ -491,12 +496,9 @@
 
             for part in partitions:
                 part_detect = 1
+                file, size, expected_crc32 = \
+                    usb_ext4load_ext4write(u_boot_console, 'ext4', x, part)
                 addr = u_boot_utils.find_ram_base(u_boot_console)
-                output = u_boot_console.run_command('crc32 %x %x' % (addr, size))
-                m = re.search('==> (.+?)', output)
-                if not m:
-                    pytest.fail('CRC32 failed')
-                expected_crc32 = m.group(1)
 
                 offset = random.randrange(128, 1024, 128)
                 output = u_boot_console.run_command(
@@ -565,15 +567,11 @@
                     addr = u_boot_utils.find_ram_base(u_boot_console)
 
                     if fs == 'fat':
-                        file, size = test_usb_fatload_fatwrite(u_boot_console)
+                        file, size, expected_crc32 = \
+                            usb_fatload_fatwrite(u_boot_console, fs, x, part)
                     elif fs == 'ext4':
-                        file, size = test_usb_ext4load_ext4write(u_boot_console)
-
-                    output = u_boot_console.run_command('crc32 %x %x' % (addr, size))
-                    m = re.search('==> (.+?)', output)
-                    if not m:
-                        pytest.fail('CRC32 failed')
-                    expected_crc32 = m.group(1)
+                        file, size, expected_crc32 = \
+                            usb_ext4load_ext4write(u_boot_console, fs, x, part)
 
                     offset = random.randrange(128, 1024, 128)
                     output = u_boot_console.run_command(