Fix argument checking to reflect new args