#!/bin/bash

show_help() {
	echo "usage: cleanup-state <pre-invariant|post-invariant>"
	echo
	echo "Runs cleanup actions of the selected class:"
	echo "   pre-invariant:"
	echo "      Clean state that is currently expected to leak from"
	echo "      any test, due to imperfections in the code."
	echo "   post-invariant:"
	echo "      Clean state that is not expected to be restored by"
	echo "      each test, that should not leak across tests."
}

main() {
	if [ $# -eq 0 ]; then
		show_help
		exit 1
	fi

	action=
	while [ $# -gt 0 ]; do
		case "$1" in
			-h|--help)
				show_help
				exit 0
				;;
			--)
				shift
				break
				;;
			pre-invariant)
				action=pre-invariant
				shift
				;;
			post-invariant)
				action=post-invariant
				shift
				;;
			-*)
				echo "cleanup-state: unsupported argument $1" >&2
				exit 1
				;;
			*)
				echo "cleanup-state: unknown action $1" >&2
				exit 1
				;;
		esac
	done

	case "$action" in
		pre-invariant)
			# If the root user has a systemd --user instance then ask it to reload.
			# This prevents tests from leaking user-session services that stay in
			# memory but are not present on disk, or have been modified on disk, as is
			# common with tests that use snaps with user services _or_ with tests that
			# cause installation of the snapd.session-agent.service unit via re-exec
			# machinery.
			#
			# This is done AHEAD of the invariant checks as it is very widespread
			# and fixing it in each test is not a priority right now.
			#
			# Note that similar treatment is not required for the "test" user as
			# correct usage of tests.session ensures that the session and all the
			# processes of the "test" user are terminated.
			if pgrep -u root --full "systemd --user"; then
				systemctl --user daemon-reload
				# Following that, check if there's a snapd.session-agent.socket and
				# if one exists stop it and then start it, ignoring errors.  If the
				# unit was removed, stopping it clears it from memory. This is
				# different from restarting the unit, which doesn't do anything if
				# the unit on disk is gone.
				if systemctl --user is-active snapd.session-agent.socket; then
					systemctl --user stop snapd.session-agent.socket
				fi
				# XXX: if there's a way to check if an unit exists but is stopped,
				# use it here to avoid starting a non-existing unit.
				systemctl --user start snapd.session-agent.socket || true

				# Do the same for root's D-Bus session bus.
				# This will stop the bus and any clients that
				# may be connected to it.
				if systemctl --user is-active dbus.socket; then
					systemctl --user stop dbus.socket
				fi
				systemctl --user start dbus.socket || true
			fi
			;;
		post-invariant)
			true
			;;
		*)
			echo "cleanup-state: unknown action $action" >&2
			exit 1
			;;
	esac
}

main "$@"
