From 7c0a5482521d369ea202712e80516c36f2b58970 Mon Sep 17 00:00:00 2001
From: Dion Hulse <contact@dd32.id.au>
Date: Fri, 11 Nov 2022 16:31:57 +1000
Subject: [PATCH] Tests: Add tests that confirm an unauthorized user cannot
 generate codes for another, but that an Administrator can, and that the codes
 are only usable by the correct user.

---
 .../class-two-factor-backup-codes-ajax.php    | 53 ++++++++++++++++++-
 1 file changed, 51 insertions(+), 2 deletions(-)

diff --git a/tests/providers/class-two-factor-backup-codes-ajax.php b/tests/providers/class-two-factor-backup-codes-ajax.php
index 063b7ebd..e8df5359 100644
--- a/tests/providers/class-two-factor-backup-codes-ajax.php
+++ b/tests/providers/class-two-factor-backup-codes-ajax.php
@@ -33,8 +33,7 @@ public function set_up() {
 	/**
 	 * Verify that the downloaded file contains the codes.
 	 *
-	 * @covers Two_Factor_Backup_Codes::generate_codes
-	 * @covers Two_Factor_Backup_Codes::validate_code
+	 * @covers Two_Factor_Backup_Codes::ajax_generate_json
 	 */
 	public function test_generate_code_and_validate_in_download_file() {
 		$this->_setRole( 'administrator' );
@@ -58,4 +57,54 @@ public function test_generate_code_and_validate_in_download_file() {
 		$this->assertTrue( $this->provider->validate_code( $user, $response->data->codes[0] ) );
 		$this->assertStringContainsString( $response->data->codes[0], $response->data->download_link );
 	}
+
+	/**
+	 * Verify that a different user cannot generate codes for another.
+	 *
+	 * @covers Two_Factor_Backup_Codes::ajax_generate_json
+	 */
+	public function test_cannot_generate_code_for_different_user() {
+		$this->_setRole( 'administrator' );
+
+		$user           = wp_get_current_user();
+		$_POST['nonce'] = wp_create_nonce( 'two-factor-backup-codes-generate-json-' . $user->ID );
+
+		// Create a new user
+		$user             = new WP_User( self::factory()->user->create() );
+		$_POST['user_id'] = $user->ID;
+
+		$this->expectException( 'WPAjaxDieStopException' );
+		$this->expectExceptionMessage( '-1' );
+		$this->_handleAjax( 'two_factor_backup_codes_generate' );
+	}
+
+	/**
+	 * Verify that an admin can create Backup codes for another user.
+	 *
+	 * @covers Two_Factor_Backup_Codes::ajax_generate_json
+	 */
+	public function test_generate_codes_for_other_users() {
+		$this->_setRole( 'administrator' );
+
+		$current_user     = wp_get_current_user();
+		$user             =  new WP_User( self::factory()->user->create() );
+		$_POST['user_id'] = $user->ID;
+		$_POST['nonce']   = wp_create_nonce( 'two-factor-backup-codes-generate-json-' . $user->ID );
+
+		try {
+			$this->_handleAjax( 'two_factor_backup_codes_generate' );
+		} catch ( WPAjaxDieContinueException $e ) {
+			unset( $e );
+		}
+
+		$this->assertStringContainsString( 'codes', $this->_last_response );
+
+		$response = json_decode( $this->_last_response );
+
+		$this->assertTrue( $response->success );
+		$this->assertNotEmpty( $response->data->codes );
+
+		$this->assertTrue( $this->provider->validate_code( $user, $response->data->codes[0] ) );
+		$this->assertFalse( $this->provider->validate_code( $current_user, $response->data->codes[0] ) );
+	}
 }