1 files changed, 46 insertions, 0 deletions
diff --git a/tests/unit/test_disk.bats b/tests/unit/test_disk.bats
index b962445..462f5d7 100644
--- a/tests/unit/test_disk.bats
+++ b/tests/unit/test_disk.bats
@@ -320,3 +320,49 @@ partition_disks_setup() {
     run disk_in_use /dev/zzz999
     [ "$status" -eq 1 ]
 }
+
+#############################
+# get_swap_partition
+#############################
+
+@test "get_swap_partition: SATA disk gets numeric suffix 3" {
+    run get_swap_partition /dev/sda
+    [ "$status" -eq 0 ]
+    [ "$output" = "/dev/sda3" ]
+}
+
+@test "get_swap_partition: NVMe disk gets p-prefixed suffix 3" {
+    run get_swap_partition /dev/nvme0n1
+    [ "$status" -eq 0 ]
+    [ "$output" = "/dev/nvme0n1p3" ]
+}
+
+#############################
+# partition_disk SWAP_SIZE layout
+#############################
+
+@test "partition_disk without SWAP_SIZE keeps the two-partition layout" {
+    partition_disks_setup
+    FILESYSTEM=zfs
+    SWAP_SIZE=""
+    partition_disk /dev/sda
+    for c in "${CALLS[@]}"; do
+        [[ "$c" != *"8200"* ]]
+    done
+}
+
+@test "partition_disk with SWAP_SIZE carves swap before root so root takes remainder" {
+    partition_disks_setup
+    FILESYSTEM=zfs
+    SWAP_SIZE=100G
+    partition_disk /dev/nvme0n1
+    local swap_idx=-1 root_idx=-1 i=0
+    for c in "${CALLS[@]}"; do
+        [[ "$c" == *"-n 3:0:+100G -t 3:8200"* ]] && swap_idx=$i
+        [[ "$c" == *"-n 2:0:0"* ]] && root_idx=$i
+        i=$((i+1))
+    done
+    [ "$swap_idx" -ge 0 ]
+    [ "$root_idx" -ge 0 ]
+    [ "$swap_idx" -lt "$root_idx" ]
+}