disk_setup: use byte string when purging the partition table (#2012)

This writes a byte string to the device instead of a string when
purging the partition table.

Essentially, this will prevent the error "a bytes-like object is
required, not 'str'" from happening.

2 files changed, 19 insertions, 2 deletions

diff --git a/cloudinit/config/cc_disk_setup.py b/cloudinit/config/cc_disk_setup.py
index 0248e7ec..4aae5530 100644
--- a/cloudinit/config/cc_disk_setup.py
+++ b/cloudinit/config/cc_disk_setup.py
@@ -657,7 +657,7 @@ def get_partition_gpt_layout(size, layout):
 def purge_disk_ptable(device):
     # wipe the first and last megabyte of a disk (or file)
     # gpt stores partition table both at front and at end.
-    null = "\0"
+    null = b"\0"
     start_len = 1024 * 1024
     end_len = 1024 * 1024
     with open(device, "rb+") as fp:
diff --git a/tests/unittests/config/test_cc_disk_setup.py b/tests/unittests/config/test_cc_disk_setup.py
index 5fa7d4cd..496ad8e1 100644
--- a/tests/unittests/config/test_cc_disk_setup.py
+++ b/tests/unittests/config/test_cc_disk_setup.py
@@ -1,6 +1,7 @@
 # This file is part of cloud-init. See LICENSE file for license information.
 
 import random
+import tempfile
 
 import pytest
 
@@ -200,6 +201,23 @@ class TestUpdateFsSetupDevices(TestCase):
         )
 
 
+class TestPurgeDisk(TestCase):
+    @mock.patch(
+        "cloudinit.config.cc_disk_setup.read_parttbl", return_value=None
+    )
+    def test_purge_disk_ptable(self, *args):
+        pseudo_device = tempfile.NamedTemporaryFile()
+
+        cc_disk_setup.purge_disk_ptable(pseudo_device.name)
+
+        with pseudo_device as f:
+            actual = f.read()
+
+        expected = b"\0" * (1024 * 1024)
+
+        self.assertEqual(expected, actual)
+
+
 @mock.patch(
     "cloudinit.config.cc_disk_setup.assert_and_settle_device",
     return_value=None,
@@ -211,7 +229,6 @@ class TestUpdateFsSetupDevices(TestCase):
 @mock.patch("cloudinit.config.cc_disk_setup.device_type", return_value=None)
 @mock.patch("cloudinit.config.cc_disk_setup.subp.subp", return_value=("", ""))
 class TestMkfsCommandHandling(CiTestCase):
-
     with_logs = True
 
     def test_with_cmd(self, subp, *args):