2 files changed, 89 insertions, 0 deletions
diff --git a/tests/unit/test_config.bats b/tests/unit/test_config.bats
index 4169c5e..554c0c7 100644
--- a/tests/unit/test_config.bats
+++ b/tests/unit/test_config.bats
@@ -302,3 +302,46 @@ EOF
     [[ "$output" == *"LUKS_PASSPHRASE"* ]]
     ! [[ "$output" == *"ZFS_PASSPHRASE"* ]]
 }
+
+#############################
+# SWAP_SIZE validation
+#############################
+
+@test "validate_config accepts empty SWAP_SIZE (default layout)" {
+    HOSTNAME=h; TIMEZONE=UTC; ROOT_PASSWORD=x
+    SELECTED_DISKS=(/dev/sda); RAID_LEVEL=""; SWAP_SIZE=""
+    run validate_config
+    [[ "$output" != *"SWAP_SIZE"* ]]
+}
+
+@test "validate_config accepts a well-formed SWAP_SIZE on a single disk" {
+    HOSTNAME=h; TIMEZONE=UTC; ROOT_PASSWORD=x
+    SELECTED_DISKS=(/dev/sda); RAID_LEVEL=""; SWAP_SIZE=100G
+    run validate_config
+    [[ "$output" != *"Invalid SWAP_SIZE"* ]]
+    [[ "$output" != *"single-disk"* ]]
+}
+
+@test "validate_config rejects SWAP_SIZE without a unit suffix" {
+    HOSTNAME=h; TIMEZONE=UTC; ROOT_PASSWORD=x
+    SELECTED_DISKS=(/dev/sda); RAID_LEVEL=""; SWAP_SIZE=100
+    run validate_config
+    [ "$status" -eq 1 ]
+    [[ "$output" == *"Invalid SWAP_SIZE"* ]]
+}
+
+@test "validate_config rejects SWAP_SIZE with a bogus unit" {
+    HOSTNAME=h; TIMEZONE=UTC; ROOT_PASSWORD=x
+    SELECTED_DISKS=(/dev/sda); RAID_LEVEL=""; SWAP_SIZE=100Q
+    run validate_config
+    [ "$status" -eq 1 ]
+    [[ "$output" == *"Invalid SWAP_SIZE"* ]]
+}
+
+@test "validate_config rejects SWAP_SIZE on a multi-disk layout" {
+    HOSTNAME=h; TIMEZONE=UTC; ROOT_PASSWORD=x
+    SELECTED_DISKS=(/dev/sda /dev/sdb); RAID_LEVEL=mirror; SWAP_SIZE=100G
+    run validate_config
+    [ "$status" -eq 1 ]
+    [[ "$output" == *"single-disk"* ]]
+}
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" ]
+}