fix test

swysocki · Nov 26, 2021 · e308574 · e308574
1 parent e297700
commit e308574
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 14 deletions.
diff --git a/pygpt_disk/table.py b/pygpt_disk/table.py
@@ -88,27 +88,37 @@ def _write_header(self, header: str = "primary") -> None:
         self._write_section(self._partition_entry_size, start_byte)
         self._write_section(self._partition_array_crc, start_byte)
 
+        # once the header is written the CRC is calculated
+        self._checksum_header(start_byte)
+
     def _write_section(self, entry: HeaderEntry, buffer_position: int):
+        """Write a GPT header entry
+
+        Args:
+            entry: HeaderEntry object
+            buffer_position: the byte offset of where the entry should be written
+        """
         # seek to section's offset
         self.disk.buffer.seek(buffer_position + entry.offset)
         if type(entry.content) != bytes:
             self.disk.buffer.write((entry.content).to_bytes(entry.length, "little"))
         else:
             self.disk.buffer.write(entry.content)
 
-    def _checksum_header(self, primary=True):
-        """Calculate the header checksum"""
-        start_byte = self.primary_header_start_byte
-        if not primary:
-            start_byte = self.backup_header_start_byte
+    def _checksum_header(self, offset: int):
+        """Calculate the header checksum
+
+        Args:
+            offset: start byte of the header we are writing (primary or backup)
+        """
         # zero field before calculating
         self._header_crc.content = 0
-        self._write_section(self._header_crc, start_byte)
+        self._write_section(self._header_crc, offset)
         # read header
-        self.disk.buffer.seek(start_byte)
+        self.disk.buffer.seek(offset)
         raw_header = self.disk.buffer.read(self._header_size.content)
         self._header_crc.content = binascii.crc32(raw_header)
-        self._write_section(self._header_crc, start_byte)
+        self._write_section(self._header_crc, offset)
 
     def update_header(self):
         pass
diff --git a/tests/test_table.py b/tests/test_table.py
@@ -30,8 +30,8 @@ def read_header(primary=True):
         assert t.disk.buffer.read(8) == SIGNATURE
         assert t.disk.buffer.read(4) == REVISION
         assert t.disk.buffer.read(4) == HEADER_SIZE
-        # header crc (zeroed until calculated)
-        assert t.disk.buffer.read(4) == b"\x00" * 4
+        # header crc ensure it is no longer zeroed
+        assert t.disk.buffer.read(4) != b"\x00" * 4
         # reserved
         assert t.disk.buffer.read(4) == b"\x00" * 4
         if primary:
@@ -72,17 +72,18 @@ def read_header(primary=True):
 
 
 def test__checksum_header(fresh_disk: disk.Disk):
+    CRC_OFFSET = 16  # offset relative to start of header
     t = table.Table(fresh_disk)
     t._write_header("primary")
-    t._checksum_header()
+    t._checksum_header(t.primary_header_start_byte)
 
     # read new checksum
-    t.disk.buffer.seek(int(PRIMARY_LBA * SECTOR_SIZE) + 16)
+    t.disk.buffer.seek(t.primary_header_start_byte + CRC_OFFSET)
     raw_crc = t.disk.buffer.read(4)
     assert raw_crc == (t._header_crc.content).to_bytes(4, "little")
 
     t._write_header("backup")
-    t._checksum_header(False)
-    t.disk.buffer.seek(int(BACKUP_LBA * SECTOR_SIZE) + 16)
+    t._checksum_header(t.backup_header_start_byte)
+    t.disk.buffer.seek(t.backup_header_start_byte + CRC_OFFSET)
     raw_crc = t.disk.buffer.read(4)
     assert raw_crc == (t._header_crc.content).to_bytes(4, "little")