#!/usr/bin/env atf-sh

. $(atf_get_srcdir)/test_env.sh
init_tests \
	setup_disk_usage \
	setup_disk_mode_none \
	setup_disk_none \
	setup_disk_func_find_disks \
	setup_disk_func_setup_partitions_dos \
	setup_disk_func_setup_partitions_gpt \
	setup_disk_func_find_efi_size \
	setup_disk_func_find_efi_size_multi \
	setup_disk_func_uuid_or_device \
	setup_disk_non_existing_block_dev \
	setup_disk_sys \
	setup_disk_sys_wifi \
	setup_disk_install_mounted_root_nvme \
	setup_disk_install_mounted_root_efi \
	setup_disk_install_mounted_root_efi_boot \
	setup_disk_install_mounted_root_bootloader_none \
	setup_disk_install_mounted_root_zfs

setup_disk_usage_body() {
	test_usage setup-disk
}

setup_disk_mode_none_body() {
	init_env
	atf_check -s exit:0 \
		setup-disk -m none
}

setup_disk_none_body() {
	init_env
	atf_check -s exit:0 \
		setup-disk none
}

fake_disk() {
	local dev="$1"
	local size="${2:-1000}" # size in MiB
	local blocksize="${3:-512}"

	mkdir -p dev \
		sys/block/$dev/device \
		sys/block/$dev/holders \
		sys/block/$dev/queue

	echo $(( $size * 1000000 / $blocksize )) > sys/block/$dev/size
	echo $(( $blocksize )) > sys/block/$dev/queue/logical_block_size
	truncate -s "${size}M" dev/"$dev"
}


fake_partition() {
	mkdir -p sys/block/$1/$1$2/holders
}

fake_raid() {
	local md="$1"
	shift
	for dev; do
		for p in sys/block/*/$dev/holders sys/block/$dev/holders; do
			if [ -d "$p" ]; then
				touch "$p"/$md
			fi
		done
	done
}

setup_disk_func_find_disks_body() {
	init_env
	fake_disk vda
	fake_disk vdb
	fake_partition vdb 1
	fake_disk sda
	fake_disk nvme0n1
	fake_partition nvme0n1 p2

	# simulate vda and vdb1 being part of md0 raid
	fake_raid md0 vda vdb1

	# simulate nvme0n1p2 being mounted
	fake_mount '/dev/nvme0n1p2 /boot vfat rw,relatime,fmask=0022 0 0'

	SETUP_DISK_TESTFUNC=find_disks USE_RAID=1 \
		atf_check -s exit:0 \
		-o match:"^ sda$" \
		-e match:"vda is part of a running raid" \
		-e match:"vdb is part of a running raid" \
		setup-disk
}

setup_disk_func_setup_partitions_dos_body() {
	atf_require_prog sfdisk
	init_env

	mkdir -p dev
	truncate -s 100M dev/vda

	SETUP_DISK_TESTFUNC=setup_partitions  \
		atf_check -s exit:0 \
		setup-disk dev/vda "0M," "10M,83,*", "10M,82", "10M,fd", ",8e"

	SETUP_DISK_TESTFUNC=find_partitions  \
		atf_check -s exit:0 \
		-o match:"dev/vda1" \
		setup-disk dev/vda "boot"

	SETUP_DISK_TESTFUNC=find_partitions  \
		atf_check -s exit:0 \
		-o match:"dev/vda1" \
		setup-disk dev/vda "linux"

	SETUP_DISK_TESTFUNC=find_partitions  \
		atf_check -s exit:0 \
		-o match:"dev/vda2" \
		setup-disk dev/vda "swap"

	SETUP_DISK_TESTFUNC=find_partitions  \
		atf_check -s exit:0 \
		-o match:"dev/vda3" \
		setup-disk dev/vda "raid"

	SETUP_DISK_TESTFUNC=find_partitions  \
		atf_check -s exit:0 \
		-o match:"dev/vda4" \
		setup-disk dev/vda "lvm"

	SETUP_DISK_TESTFUNC=find_nth_non_boot_parts \
		atf_check -s exit:0 \
		setup-disk 1 "83" dev/vda
}

setup_disk_func_setup_partitions_gpt_body() {
	atf_require_prog sfdisk
	init_env

	mkdir -p dev
	truncate -s 200M dev/vda

	swap=0657FD6D-A4AB-43C4-84E5-0933C84B4F4F
	linux=0FC63DAF-8483-4772-8E79-3D69D8477DE4
	raid=A19D880F-05FC-4D3B-A006-743F0F84911E
	lvm=E6D6D379-F507-44C2-A23C-238F2A3DF928
	prep=9E1A2d38-C612-4316-AA26-8B49521E5A8B
	esp=C12A7328-F81F-11D2-BA4B-00A0C93EC93B

	SETUP_DISK_TESTFUNC=setup_partitions DISKLABEL=gpt \
		atf_check -s exit:0 \
		setup-disk dev/vda "0M," "30M,$esp" "30M,$swap" "size=30M,type=$linux,attrs=LegacyBIOSBootable" "30M,$raid" "30M,$lvm" "size=10M,type=linux"

	sfdisk -d dev/vda

	SETUP_DISK_TESTFUNC=find_partitions DISKLABEL=gpt USE_EFI=1 \
		atf_check -s exit:0 \
		-o match:"dev/vda1" \
		setup-disk dev/vda "boot"

	# zsh needs this
	unset USE_EFI

	SETUP_DISK_TESTFUNC=find_partitions DISKLABEL=gpt \
		atf_check -s exit:0 \
		-o match:"dev/vda2" \
		setup-disk dev/vda "swap"

	SETUP_DISK_TESTFUNC=find_partitions DISKLABEL=gpt \
		atf_check -s exit:0 \
		-o match:"dev/vda3" \
		setup-disk dev/vda "linux"

	SETUP_DISK_TESTFUNC=find_partitions DISKLABEL=gpt \
		atf_check -s exit:0 \
		-o match:"dev/vda4" \
		setup-disk dev/vda "raid"

	SETUP_DISK_TESTFUNC=find_partitions DISKLABEL=gpt \
		atf_check -s exit:0 \
		-o match:"dev/vda5" \
		setup-disk dev/vda "lvm"

	SETUP_DISK_TESTFUNC=find_nth_non_boot_parts DISKLABEL=gpt \
		atf_check -s exit:0 \
		-o match:"dev/vda6" \
		setup-disk 1 "$linux" dev/vda
}

setup_disk_func_find_efi_size_body() {
	init_env
	# size:blocksize:expected
	for i in 10000::512 4000::264 1000::160 \
			10000:4096:512 4000:4096:264 1000:4096:264; do

		local diskdata=${i%:*}
		local disksize=${diskdata%:*}
		local blocksize=${diskdata#*:}
		local expected=${i##*:}

		fake_disk vda $disksize $blocksize
		SETUP_DISK_TESTFUNC=find_efi_size \
			atf_check -s exit:0 \
			-o match:"^$expected$" \
			setup-disk /dev/vda
	done
}

setup_disk_func_find_efi_size_multi_body() {
	init_env
	# size:blocksize:expected
	for i in 10000::512 4000::264 1000::160 \
			10000:4096:512 4000:4096:264 1000:4096:264; do

		local diskdata=${i%:*}
		local disksize=${diskdata%:*}
		local blocksize=${diskdata#*:}
		local expected=${i##*:}

		fake_disk vda $disksize $blocksize
		fake_disk vdb $disksize $blocksize
		SETUP_DISK_TESTFUNC=find_efi_size \
			atf_check -s exit:0 \
			-o match:"^$expected$" \
			setup-disk /dev/vda /dev/vdb
	done
}

setup_disk_func_uuid_or_device_body() {
	init_env
	fake_bin blkid <<-EOF
		#!/bin/sh
		cat <<-EOF2
			/dev/sda3: UUID="4ec0992a-bf5a-4d04-bc49-06c131ae901b" BLOCK_SIZE="4096" TYPE="ext4" PARTUUID="0ce794d3-03"
			/dev/sda1: UUID="1db12961-31ce-420f-94f0-559a1f6592ff" BLOCK_SIZE="1024" TYPE="ext4" PARTUUID="0ce794d3-01"
			/dev/sda2: UUID="a5fbbd70-94f6-4d03-947c-e3afb91a32c8" TYPE="swap" PARTUUID="0ce794d3-02"
			/dev/md0: UUID="86a00eb9-9b5a-4d73-812c-8fc958673a7c" TYPE="LVM2_member"
		EOF2
	EOF

	SETUP_DISK_TESTFUNC=uuid_or_device \
		atf_check -s exit:0 \
			-o match:"/dev/vda1" \
			setup-disk /dev/vda1

	SETUP_DISK_TESTFUNC=uuid_or_device \
		atf_check -s exit:0 \
			-o not-match:"UUID=" \
			-o match:"/dev/md0" \
			setup-disk /dev/md0

	SETUP_DISK_TESTFUNC=uuid_or_device \
		atf_check -s exit:0 \
			-o match:"UUID=1db12961-31ce-420f-94f0-559a1f6592ff" \
			setup-disk /dev/sda1

	SETUP_DISK_TESTFUNC=uuid_or_device \
	NO_DISK_UUID="/dev/sda1 /dev/sda3" \
		atf_check -s exit:0 \
			-o not-match:"UUID" \
			-o match:"/dev/sda1" \
			setup-disk /dev/sda1

}

setup_disk_non_existing_block_dev_body() {
	init_env
	atf_check -s not-exit:0 \
		-e match:"/dev/vda is not a block device suitable for partitioning" \
		-o match:"swapoff" \
		setup-disk -m sys /dev/vda
}

setup_disk_sys_body() {
	atf_require_prog sfdisk
	init_env
	fake_disk vda
	mkdir proc etc
	touch proc/mounts etc/fstab
	ERASE_DISKS=/dev/vda atf_check \
		-o match:"Installation is complete" \
		setup-disk -m sys /dev/vda

	atf_check \
		-o match:"^ext4$"\
		awk '$2 == "/" {print $3}' mnt/etc/fstab
	atf_check \
		-o match:"^ext4$"\
		awk '$2 == "/boot" {print $3}' mnt/etc/fstab
	atf_check \
		-o inline:"none\n"\
		awk '$3 == "swap" {print $2}' mnt/etc/fstab
}

setup_disk_sys_wifi_body() {
	atf_require_prog sfdisk
	init_env
	fake_disk vda
	mkdir -p proc etc sys/class/net/wlan0/wireless
	touch proc/mounts etc/fstab
	ERASE_DISKS=/dev/vda atf_check \
		-o match:"apk add.*wireless-regdb" \
		setup-disk -m sys /dev/vda
}


setup_disk_install_mounted_root_nvme_body() {
	init_env
	mkdir -p target/boot/efi \
		sys/firmware/efi

	# simulate nvme0n1p2 being mounted
	fake_mount "/dev/nvme0n1p2 $PWD/target ext4 rw,noatime,data=ordered 0 0"
	fake_mount "/dev/nvme0n1p3 $PWD/target/boot/efi vfat rw,relatime,fmask=0022 0 0"

	atf_check -s exit:0 \
		-o match:"Installing system on /dev/nvme0n1p2" \
		-o match:"grub-install .* --efi-directory=$PWD/target/boot/efi .* --boot-directory=$PWD/target/boot" \
		-o match:"install -D $PWD/target/boot/efi/EFI/alpine/grub.*efi $PWD/target/boot/efi/EFI/boot/bootaa64.efi" \
		setup-disk target

	atf_check \
		-o match:"/dev/nvme0n1p2	/	ext4	rw,noatime,data=ordered 0 1" \
		-o match:"/dev/nvme0n1p3	/boot/efi	vfat	rw,relatime,fmask=0022 0 2" \
		cat target/etc/fstab

	atf_check -o match:"GRUB_CMDLINE_LINUX_DEFAULT=.*ext4,nvme" \
		cat target/etc/default/grub

}

setup_disk_install_mounted_root_efi_body() {
	init_env
	for fs in vfat msdos fat; do
		rm -rf target sys etc proc
		mkdir -p target/boot/efi \
			sys/firmware/efi \

		fake_mount "/dev/vda2 $PWD/target ext4 rw,noatime,data=ordered 0 0"
		fake_mount "/dev/vda1 $PWD/target/boot/efi $fs rw,relatime,fmask=0022 0 0"

		atf_check -s exit:0 \
			-o match:"Installing system on /dev/vda2" \
			-o match:"grub-install .* --efi-directory=$PWD/target/boot/efi .* --boot-directory=$PWD/target/boot" \
			-o match:"install -D $PWD/target/boot/efi/EFI/alpine/grub.*efi $PWD/target/boot/efi/EFI/boot/bootaa64.efi" \
			setup-disk -v target

		atf_check \
			-o match:"/dev/vda2	/	ext4	rw,noatime,data=ordered 0 1" \
			-o match:"/dev/vda1	/boot/efi	$fs	rw,relatime,fmask=0022 0 2" \
			cat target/etc/fstab

		atf_check -o match:"GRUB_CMDLINE_LINUX_DEFAULT=.*ext4" \
			cat target/etc/default/grub
	done
}

# test case when /boot is the efi vfat partition
setup_disk_install_mounted_root_efi_boot_body() {
	init_env
	mkdir -p target/boot \
		sys/firmware/efi

	fake_mount "/dev/vda2 $PWD/target ext4 rw,noatime,data=ordered 0 0"
	fake_mount "/dev/vda1 $PWD/target/boot vfat rw,relatime,fmask=0022 0 0"

	atf_check -s exit:0 \
		-o match:"Installing system on /dev/vda2" \
		-o match:"grub-install .* --efi-directory=$PWD/target/boot .* --boot-directory=$PWD/target/boot" \
		-o match:"install -D $PWD/target/boot/EFI/alpine/grub.*efi $PWD/target/boot/EFI/boot/bootaa64.efi" \
		setup-disk -v target

	atf_check \
		-o match:"/dev/vda2	/	ext4	rw,noatime,data=ordered 0 1" \
		-o match:"/dev/vda1	/boot	vfat	rw,relatime,fmask=0022 0 2" \
		cat target/etc/fstab

	atf_check -o match:"GRUB_CMDLINE_LINUX_DEFAULT=.*ext4" \
		cat target/etc/default/grub
}

setup_disk_install_mounted_root_bootloader_none_body() {
	init_env
	mkdir -p target/boot \
		sys/firmware/efi

	fake_mount "/dev/vda2 $PWD/target virtiofs rw,noatime,data=ordered 0 0"

	BOOTLOADER=none atf_check -s exit:0 \
		-o not-match:"grub" \
		setup-disk -v target

	if [ -f target/etc/default/grub ]; then
		atf_fail "etc/default/grub should not exist"
	fi
}

setup_disk_install_mounted_root_zfs_body() {
	init_env
	mkdir -p target/boot \
		sys/firmware/efi

	fake_mount "rpool/root-a $PWD/target zfs rw,noatime,data=ordered 0 0"

	atf_check -s exit:0 \
		-o match:"grub" \
		-e not-match:"zfs is not supported" \
		setup-disk -v target
	atf_check \
		-o match:"GRUB_CMDLINE_LINUX_DEFAULT=.*zfs" \
		-o match:"GRUB_FS=zfs" \
		cat target/etc/default/grub
	atf_check \
		-o match:'features=.*nvme' \
		-o match:'features=.*zfs' \
		-o match:'features=.*keymap' \
		cat target/etc/mkinitfs/mkinitfs.conf
}

