From fa8dded7795f62e381c999114ee09e69ac3882b5 Mon Sep 17 00:00:00 2001
From: Echo G <ChaosExAnima@users.noreply.github.com>
Date: Thu, 26 Mar 2020 17:00:06 -0400
Subject: [PATCH] Adds unit tests for added methods.

---
 class.jetpack.php                        |  2 +-
 tests/php/general/test_class.jetpack.php | 74 ++++++++++++++++++++++++
 2 files changed, 75 insertions(+), 1 deletion(-)

diff --git a/class.jetpack.php b/class.jetpack.php
index 1e70022147908..e277654c00710 100644
--- a/class.jetpack.php
+++ b/class.jetpack.php
@@ -384,7 +384,7 @@ class Jetpack {
 	 * @var string
 	 * @since 8.4.0
 	 */
-	protected static $jetpack_redirect_login = 'jetpack_connect_login_redirect';
+	public static $jetpack_redirect_login = 'jetpack_connect_login_redirect';
 
 	/**
 	 * Holds the singleton instance of this class
diff --git a/tests/php/general/test_class.jetpack.php b/tests/php/general/test_class.jetpack.php
index c646a5570b6a8..7dffaedfb1cae 100644
--- a/tests/php/general/test_class.jetpack.php
+++ b/tests/php/general/test_class.jetpack.php
@@ -1208,4 +1208,78 @@ public function partner_code_provider() {
 		);
 	}
 
+	/**
+	 * Tests login URL only adds redirect param when redirect param is in original request.
+	 *
+	 * @since 8.4.0
+	 * @return void
+	 */
+	public function test_login_url_add_redirect() {
+		$login_url = wp_login_url( '/wp-admin' );
+		$this->assertStringNotContainsString( $login_url, Jetpack::$jetpack_redirect_login );
+
+		$login_url = wp_login_url( '/wp-admin?' . Jetpack::$jetpack_redirect_login . '=true' );
+		parse_str( wp_parse_url( $login_url, PHP_URL_QUERY ), $login_parts );
+		$this->assertArraySubset( array( Jetpack::$jetpack_redirect_login => 'true' ), $login_parts, true );
+	}
+
+	/**
+	 * Tests login redirect sending users to Calypso when redirect param is set.
+	 *
+	 * @since 8.4.0
+	 * @return void
+	 */
+	public function test_login_init_redirect() {
+		tests_add_filter(
+			'wp_redirect',
+			function( $location ) {
+				$expected_location = add_query_arg(
+					array(
+						'forceInstall' => 1,
+						'url'          => rawurlencode( get_site_url() ),
+					),
+					'https://wordpress.com/jetpack/connect'
+				);
+				$this->assertEquals( $location, $expected_location );
+				throw new Exception(); // Cause an exception, as we don't want to run exit.
+			}
+		);
+
+		// Remove core filters that add headers.
+		remove_filter( 'login_init', 'wp_admin_headers' );
+		remove_filter( 'login_init', 'send_frame_options_header' );
+
+		// Run it once and no exception is thrown.
+		do_action( 'login_init' );
+
+		$this->expectException( Exception::class );
+		$_GET[ Jetpack::$jetpack_redirect_login ] = 'true';
+		do_action( 'login_init' ); // Now expect an exception.
+	}
+
+	/**
+	 * Tests getting the correct Calypso host.
+	 *
+	 * @since 8.4.0
+	 * @return void
+	 */
+	public function test_get_calypso_host() {
+		// No env.
+		$this->assertEquals( 'https://wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'development';
+		$this->assertEquals( 'http://calypso.localhost:3000/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'wpcalypso';
+		$this->assertEquals( 'https://wpcalypso.wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'horizon';
+		$this->assertEquals( 'https://horizon.wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'stage';
+		$this->assertEquals( 'https://wordpress.com/', Jetpack::get_calypso_host() );
+
+		$_GET['calypso_env'] = 'production';
+		$this->assertEquals( 'https://wordpress.com/', Jetpack::get_calypso_host() );
+	}
 } // end class