Commit 1f39d1b779 for woocommerce

commit 1f39d1b77903002ad0ae3629314e70a5d2250b8e
Author: Hannah Tinkler <hannah.tinkler@gmail.com>
Date:   Wed Jan 14 17:11:05 2026 +0000

    Adds controller boilerplate for the push token registration endpoint (#62751)

    * Adds controller foundation and tests.

diff --git a/plugins/woocommerce/src/Internal/PushNotifications/Controllers/PushTokenRestController.php b/plugins/woocommerce/src/Internal/PushNotifications/Controllers/PushTokenRestController.php
new file mode 100644
index 0000000000..5fd35c3c63
--- /dev/null
+++ b/plugins/woocommerce/src/Internal/PushNotifications/Controllers/PushTokenRestController.php
@@ -0,0 +1,227 @@
+<?php
+
+declare( strict_types = 1 );
+
+namespace Automattic\WooCommerce\Internal\PushNotifications\Controllers;
+
+defined( 'ABSPATH' ) || exit;
+
+use Automattic\WooCommerce\Internal\PushNotifications\DataStores\PushTokensDataStore;
+use Automattic\WooCommerce\Internal\PushNotifications\Entities\PushToken;
+use Automattic\WooCommerce\Internal\PushNotifications\Exceptions\PushTokenNotFoundException;
+use Automattic\WooCommerce\Internal\PushNotifications\PushNotifications;
+use Automattic\WooCommerce\Internal\RestApiControllerBase;
+use InvalidArgumentException;
+use Exception;
+use WP_REST_Request;
+use WP_Error;
+use WP_Http;
+
+/**
+ * Controller for the REST endpoints associated with push notification device
+ * tokens.
+ *
+ * @since 10.6.0
+ */
+class PushTokenRestController extends RestApiControllerBase {
+	/**
+	 * The root namespace for the JSON REST API endpoints.
+	 *
+	 * @var string
+	 */
+	protected string $route_namespace = 'wc-push-notifications';
+
+	/**
+	 * The REST base for the endpoints URL.
+	 *
+	 * @var string
+	 */
+	protected string $rest_base = 'push-tokens';
+
+	/**
+	 * Get the WooCommerce REST API namespace for the class.
+	 *
+	 * @since 10.6.0
+	 *
+	 * @return string
+	 */
+	protected function get_rest_api_namespace(): string {
+		return $this->route_namespace;
+	}
+
+	/**
+	 * Register the REST API endpoints handled by this controller.
+	 *
+	 * @since 10.6.0
+	 *
+	 * @return void
+	 */
+	public function register_routes(): void {
+		// Routes will be registered here, can't omit method due to parent class
+		// constraints.
+	}
+
+	/**
+	 * Get the schema for the POST endpoint.
+	 *
+	 * @since 10.6.0
+	 *
+	 * @return array[]
+	 */
+	public function get_schema(): array {
+		return array_merge(
+			$this->get_base_schema(),
+			array(
+				'title'      => PushToken::POST_TYPE,
+				'properties' => array_map(
+					fn ( $item ) => array_intersect_key(
+						$item,
+						array(
+							'description' => null,
+							'type'        => null,
+							'enum'        => null,
+							'minimum'     => null,
+							'default'     => null,
+							'required'    => null,
+						)
+					),
+					$this->get_args()
+				),
+			)
+		);
+	}
+
+	/**
+	 * Checks user is authorized to access this endpoint.
+	 *
+	 * @since 10.6.0
+	 *
+	 * @param WP_REST_Request $request The request object.
+	 * @phpstan-param WP_REST_Request<array<string, mixed>> $request
+	 * @return bool|WP_Error
+	 */
+	public function authorize( WP_REST_Request $request ) {
+		if (
+			! get_current_user_id()
+			|| ! wc_get_container()->get( PushNotifications::class )->should_be_enabled()
+		) {
+			return false;
+		}
+
+		$has_valid_role = array_reduce(
+			PushNotifications::ROLES_WITH_PUSH_NOTIFICATIONS_ENABLED,
+			fn ( $carry, $role ) => $this->check_permission( $request, $role ) === true ? true : $carry,
+			false
+		);
+
+		if ( ! $has_valid_role ) {
+			return false;
+		}
+
+		if ( $request->has_param( 'id' ) ) {
+			$push_token = new PushToken();
+			$push_token->set_id( (int) $request->get_param( 'id' ) );
+
+			try {
+				wc_get_container()->get( PushTokensDataStore::class )->read( $push_token );
+			} catch ( Exception $e ) {
+				return $this->convert_exception_to_wp_error( $e );
+			}
+
+			if ( $push_token->get_user_id() !== get_current_user_id() ) {
+				return new WP_Error(
+					'rest_invalid_push_token',
+					'Push token could not be found.',
+					array( 'status' => WP_Http::NOT_FOUND )
+				);
+			}
+		}
+
+		return true;
+	}
+
+	/**
+	 * Converts an exception to an instance of WP_Error.
+	 *
+	 * @since 10.6.0
+	 *
+	 * @param Exception $e The exception to convert.
+	 * @return WP_Error
+	 */
+	private function convert_exception_to_wp_error( Exception $e ): WP_Error {
+		$exception_class = get_class( $e );
+
+		$slugs = array(
+			PushTokenNotFoundException::class => 'rest_invalid_push_token',
+			InvalidArgumentException::class   => 'rest_invalid_argument',
+		);
+
+		$statuses = array(
+			PushTokenNotFoundException::class => WP_Http::NOT_FOUND,
+			InvalidArgumentException::class   => WP_Http::BAD_REQUEST,
+		);
+
+		$slug   = $slugs[ $exception_class ] ?? 'rest_internal_error';
+		$status = $statuses[ $exception_class ] ?? WP_Http::INTERNAL_SERVER_ERROR;
+
+		return new WP_Error( $slug, $e->getMessage(), array( 'status' => $status ) );
+	}
+
+	/**
+	 * Get the accepted arguments for the POST request.
+	 *
+	 * @since 10.6.0
+	 *
+	 * @param string $context The context to return args for.
+	 * @return array
+	 */
+	private function get_args( ?string $context = null ): array {
+		$args = array(
+			'id'          => array(
+				'description'       => __( 'Push Token ID', 'woocommerce' ),
+				'type'              => 'integer',
+				'required'          => true,
+				'context'           => array( 'delete' ),
+				'minimum'           => 1,
+				'sanitize_callback' => 'absint',
+			),
+			'origin'      => array(
+				'description' => __( 'Origin', 'woocommerce' ),
+				'type'        => 'string',
+				'required'    => true,
+				'context'     => array( 'create' ),
+				'enum'        => PushToken::ORIGINS,
+			),
+			'device_uuid' => array(
+				'description'       => __( 'Device UUID', 'woocommerce' ),
+				'default'           => '',
+				'type'              => 'string',
+				'context'           => array( 'create' ),
+				'sanitize_callback' => 'sanitize_text_field',
+			),
+			'platform'    => array(
+				'description' => __( 'Platform', 'woocommerce' ),
+				'type'        => 'string',
+				'required'    => true,
+				'context'     => array( 'create' ),
+				'enum'        => PushToken::PLATFORMS,
+			),
+			'token'       => array(
+				'description'       => __( 'Push Token', 'woocommerce' ),
+				'type'              => 'string',
+				'required'          => true,
+				'context'           => array( 'create' ),
+				'sanitize_callback' => 'wp_unslash',
+			),
+		);
+
+		if ( $context ) {
+			$args = array_filter(
+				$args,
+				fn ( $arg ) => in_array( $context, $arg['context'], true )
+			);
+		}
+
+		return $args;
+	}
+}
diff --git a/plugins/woocommerce/src/Internal/PushNotifications/PushNotifications.php b/plugins/woocommerce/src/Internal/PushNotifications/PushNotifications.php
index d89b87fa80..bc4d75aa7a 100644
--- a/plugins/woocommerce/src/Internal/PushNotifications/PushNotifications.php
+++ b/plugins/woocommerce/src/Internal/PushNotifications/PushNotifications.php
@@ -7,6 +7,7 @@ namespace Automattic\WooCommerce\Internal\PushNotifications;
 defined( 'ABSPATH' ) || exit;

 use Automattic\Jetpack\Connection\Manager as JetpackConnectionManager;
+use Automattic\WooCommerce\Internal\PushNotifications\Controllers\PushTokenRestController;
 use Automattic\WooCommerce\Internal\PushNotifications\Entities\PushToken;
 use Automattic\WooCommerce\Proxies\LegacyProxy;
 use Automattic\WooCommerce\Utilities\FeaturesUtil;
@@ -44,18 +45,31 @@ class PushNotifications {
 	private ?bool $enabled = null;

 	/**
-	 * Loads the push notifications class.
+	 * Registers initialisation tasks to the `init` hook.
 	 *
 	 * @return void
 	 *
 	 * @since 10.4.0
 	 */
 	public function register(): void {
+		add_action( 'init', array( $this, 'on_init' ) );
+	}
+
+	/**
+	 * Loads the push notifications class.
+	 *
+	 * @return void
+	 *
+	 * @since 10.6.0
+	 */
+	public function on_init(): void {
 		if ( ! $this->should_be_enabled() ) {
 			return;
 		}

-		add_action( 'init', array( $this, 'register_post_types' ) );
+		$this->register_post_types();
+
+		wc_get_container()->get( PushTokenRestController::class )->register();

 		// Library endpoints and scheduled tasks will be registered here.
 	}
diff --git a/plugins/woocommerce/tests/php/src/Internal/PushNotifications/Controllers/PushTokenRestControllerTest.php b/plugins/woocommerce/tests/php/src/Internal/PushNotifications/Controllers/PushTokenRestControllerTest.php
new file mode 100644
index 0000000000..d5542e22f6
--- /dev/null
+++ b/plugins/woocommerce/tests/php/src/Internal/PushNotifications/Controllers/PushTokenRestControllerTest.php
@@ -0,0 +1,323 @@
+<?php
+
+declare( strict_types = 1 );
+
+namespace Automattic\WooCommerce\Tests\Internal\PushNotifications\Controllers;
+
+use Automattic\Jetpack\Connection\Manager as JetpackConnectionManager;
+use Automattic\WooCommerce\Internal\Features\FeaturesController;
+use Automattic\WooCommerce\Internal\PushNotifications\Controllers\PushTokenRestController;
+use Automattic\WooCommerce\Internal\PushNotifications\DataStores\PushTokensDataStore;
+use Automattic\WooCommerce\Internal\PushNotifications\Entities\PushToken;
+use Automattic\WooCommerce\Internal\PushNotifications\PushNotifications;
+use Automattic\WooCommerce\Proxies\LegacyProxy;
+use PHPUnit\Framework\MockObject\MockObject;
+use ReflectionClass;
+use WC_REST_Unit_Test_Case;
+use WP_Error;
+use WP_Http;
+use WP_REST_Request;
+
+/**
+ * Tests for the PushTokenRestController class.
+ *
+ * @package WooCommerce\Tests\PushNotifications
+ */
+class PushTokenRestControllerTest extends WC_REST_Unit_Test_Case {
+	/**
+	 * User ID for testing.
+	 *
+	 * @var int
+	 */
+	private $user_id;
+
+	/**
+	 * @var JetpackConnectionManager|MockObject
+	 */
+	private $jetpack_connection_manager_mock;
+
+	/**
+	 * @var FeaturesController|MockObject
+	 */
+	private $features_controller_mock;
+
+	/**
+	 * @var PushTokenRestController
+	 */
+	private $controller;
+
+	/**
+	 * Set up test.
+	 */
+	public function setUp(): void {
+		parent::setUp();
+
+		$this->set_up_features_controller_mock();
+		$this->reset_push_notifications_cache();
+
+		$this->controller = new PushTokenRestController();
+		$this->user_id    = $this->factory->user->create( array( 'role' => 'shop_manager' ) );
+	}
+
+	/**
+	 * Tear down test.
+	 */
+	public function tearDown(): void {
+		wp_set_current_user( 0 );
+
+		$this->reset_container_replacements();
+		wc_get_container()->reset_all_resolved();
+
+		parent::tearDown();
+	}
+
+	/**
+	 * @testdox Test authorize returns false when no user is logged in.
+	 */
+	public function test_authorize_returns_false_without_authentication() {
+		$this->mock_jetpack_connection_manager_is_connected( true );
+
+		$request = new WP_REST_Request( 'POST', '/wc-push-notifications/push-tokens' );
+		$result  = $this->controller->authorize( $request );
+
+		$this->assertFalse( $result );
+	}
+
+	/**
+	 * @testdox Test authorize returns false when push notifications are disabled.
+	 */
+	public function test_authorize_returns_false_when_push_notifications_disabled() {
+		wp_set_current_user( $this->user_id );
+
+		$this->mock_jetpack_connection_manager_is_connected( false );
+
+		$request = new WP_REST_Request( 'POST', '/wc-push-notifications/push-tokens' );
+		$result  = $this->controller->authorize( $request );
+
+		$this->assertFalse( $result );
+	}
+
+	/**
+	 * @testdox Test authorize returns false when user doesn't have required role.
+	 */
+	public function test_authorize_returns_false_without_required_role() {
+		$customer_id = $this->factory->user->create( array( 'role' => 'customer' ) );
+		wp_set_current_user( $customer_id );
+
+		$this->mock_jetpack_connection_manager_is_connected( true );
+
+		$request = new WP_REST_Request( 'POST', '/wc-push-notifications/push-tokens' );
+		$result  = $this->controller->authorize( $request );
+
+		$this->assertFalse( $result );
+	}
+
+	/**
+	 * @testdox Test authorize returns true for shop_manager role.
+	 */
+	public function test_authorize_returns_true_for_shop_manager() {
+		wp_set_current_user( $this->user_id );
+
+		$this->mock_jetpack_connection_manager_is_connected( true );
+
+		$request = new WP_REST_Request( 'POST', '/wc-push-notifications/push-tokens' );
+		$result  = $this->controller->authorize( $request );
+
+		$this->assertTrue( $result );
+	}
+
+	/**
+	 * @testdox Test authorize returns true for administrator role.
+	 */
+	public function test_authorize_returns_true_for_administrator() {
+		$admin_id = $this->factory->user->create( array( 'role' => 'administrator' ) );
+		wp_set_current_user( $admin_id );
+
+		$this->mock_jetpack_connection_manager_is_connected( true );
+
+		$request = new WP_REST_Request( 'POST', '/wc-push-notifications/push-tokens' );
+		$result  = $this->controller->authorize( $request );
+
+		$this->assertTrue( $result );
+	}
+
+	/**
+	 * @testdox Test authorize returns error when token ID doesn't exist.
+	 */
+	public function test_authorize_returns_error_when_token_not_found() {
+		wp_set_current_user( $this->user_id );
+
+		$this->mock_jetpack_connection_manager_is_connected( true );
+
+		$request = new WP_REST_Request( 'DELETE', '/wc-push-notifications/push-tokens/999999' );
+		$request->set_param( 'id', 999999 );
+		$result = $this->controller->authorize( $request );
+
+		$this->assertInstanceOf( WP_Error::class, $result );
+		$this->assertEquals( 'rest_invalid_push_token', $result->get_error_code() );
+		$this->assertEquals( WP_Http::NOT_FOUND, $result->get_error_data()['status'] );
+	}
+
+	/**
+	 * @testdox Test authorize returns error when token belongs to another user.
+	 */
+	public function test_authorize_returns_error_when_token_belongs_to_another_user() {
+		/**
+		 * Create a token for another shop manager.
+		 */
+		$other_user_id = $this->factory->user->create( array( 'role' => 'shop_manager' ) );
+
+		$push_token = new PushToken();
+		$push_token->set_user_id( $other_user_id );
+		$push_token->set_token( str_repeat( 'a', 64 ) );
+		$push_token->set_platform( PushToken::PLATFORM_APPLE );
+		$push_token->set_device_uuid( 'device-other-user' );
+		$push_token->set_origin( PushToken::ORIGIN_WOOCOMMERCE_IOS );
+
+		$data_store = wc_get_container()->get( PushTokensDataStore::class );
+		$data_store->create( $push_token );
+		$token_id = $push_token->get_id();
+
+		/**
+		 * Try to authorize as a different user.
+		 */
+		wp_set_current_user( $this->user_id );
+
+		$this->mock_jetpack_connection_manager_is_connected( true );
+
+		$request = new WP_REST_Request( 'DELETE', '/wc-push-notifications/push-tokens/' . $token_id );
+		$request->set_param( 'id', $token_id );
+
+		$result = $this->controller->authorize( $request );
+
+		$this->assertInstanceOf( WP_Error::class, $result );
+		$this->assertEquals( 'rest_invalid_push_token', $result->get_error_code() );
+		$this->assertEquals( WP_Http::NOT_FOUND, $result->get_error_data()['status'] );
+	}
+
+	/**
+	 * @testdox Test authorize returns true when token belongs to current user.
+	 */
+	public function test_authorize_returns_true_when_token_belongs_to_current_user() {
+		wp_set_current_user( $this->user_id );
+
+		$this->mock_jetpack_connection_manager_is_connected( true );
+
+		$push_token = new PushToken();
+		$push_token->set_user_id( $this->user_id );
+		$push_token->set_token( str_repeat( 'a', 64 ) );
+		$push_token->set_platform( PushToken::PLATFORM_APPLE );
+		$push_token->set_device_uuid( 'device-current-user' );
+		$push_token->set_origin( PushToken::ORIGIN_WOOCOMMERCE_IOS );
+
+		$data_store = wc_get_container()->get( PushTokensDataStore::class );
+		$data_store->create( $push_token );
+		$token_id = $push_token->get_id();
+
+		$request = new WP_REST_Request( 'DELETE', '/wc-push-notifications/push-tokens/' . $token_id );
+		$request->set_param( 'id', $token_id );
+
+		$result = $this->controller->authorize( $request );
+
+		$this->assertTrue( $result );
+	}
+
+	/**
+	 * @testdox Test the schema is correctly formatted.
+	 */
+	public function test_get_schema_returns_correct_structure() {
+		$schema = $this->controller->get_schema();
+
+		$this->assertArrayHasKey( 'title', $schema );
+		$this->assertArrayHasKey( 'properties', $schema );
+		$this->assertEquals( PushToken::POST_TYPE, $schema['title'] );
+
+		$this->assertArrayHasKey( 'token', $schema['properties'] );
+		$this->assertArrayHasKey( 'platform', $schema['properties'] );
+		$this->assertArrayHasKey( 'device_uuid', $schema['properties'] );
+		$this->assertArrayHasKey( 'origin', $schema['properties'] );
+		$this->assertArrayHasKey( 'enum', $schema['properties']['platform'] );
+		$this->assertArrayHasKey( 'enum', $schema['properties']['origin'] );
+
+		$this->assertArrayNotHasKey( 'validate_callback', $schema['properties']['token'] );
+		$this->assertArrayNotHasKey( 'validate_callback', $schema['properties']['platform'] );
+		$this->assertArrayNotHasKey( 'validate_callback', $schema['properties']['device_uuid'] );
+		$this->assertArrayNotHasKey( 'validate_callback', $schema['properties']['origin'] );
+
+		$this->assertEquals( 'string', $schema['properties']['token']['type'] );
+		$this->assertEquals( 'string', $schema['properties']['platform']['type'] );
+		$this->assertEquals( 'string', $schema['properties']['device_uuid']['type'] );
+		$this->assertEquals( 'string', $schema['properties']['origin']['type'] );
+
+		$this->assertEquals(
+			PushToken::PLATFORMS,
+			$schema['properties']['platform']['enum']
+		);
+
+		$this->assertEquals(
+			PushToken::ORIGINS,
+			$schema['properties']['origin']['enum']
+		);
+	}
+
+	/**
+	 * Sets up the Jetpack connection manager mocking, and ensures the
+	 * PushNotifications class state is reset so `should_be_enabled` calculates
+	 * this from scratch.
+	 *
+	 * @param bool $is_connected Whether the manager should report Jetpack is
+	 * connected or not.
+	 */
+	private function mock_jetpack_connection_manager_is_connected( bool $is_connected = true ) {
+		$this->jetpack_connection_manager_mock = $this
+			->getMockBuilder( JetpackConnectionManager::class )
+			->disableOriginalConstructor()
+			->onlyMethods( array( 'is_connected' ) )
+			->getMock();
+
+		wc_get_container()->get( LegacyProxy::class )->register_class_mocks(
+			array( JetpackConnectionManager::class => $this->jetpack_connection_manager_mock )
+		);
+
+		$this->jetpack_connection_manager_mock
+			->expects( $this->any() )
+			->method( 'is_connected' )
+			->willReturn( $is_connected );
+
+		$this->reset_push_notifications_cache();
+	}
+
+	/**
+	 * Sets up the FeaturesController mock to enable push_notifications feature.
+	 */
+	private function set_up_features_controller_mock() {
+		$this->features_controller_mock = $this
+			->getMockBuilder( FeaturesController::class )
+			->disableOriginalConstructor()
+			->onlyMethods( array( 'feature_is_enabled' ) )
+			->getMock();
+
+		$this->features_controller_mock
+			->method( 'feature_is_enabled' )
+			->willReturnCallback(
+				function ( $feature_id ) {
+					return PushNotifications::FEATURE_NAME === $feature_id;
+				}
+			);
+
+		wc_get_container()->replace( FeaturesController::class, $this->features_controller_mock );
+	}
+
+	/**
+	 * Resets the cached enablement state on the container's PushNotifications
+	 * instance.
+	 */
+	private function reset_push_notifications_cache() {
+		$push_notifications = wc_get_container()->get( PushNotifications::class );
+		$reflection         = new ReflectionClass( $push_notifications );
+		$property           = $reflection->getProperty( 'enabled' );
+
+		$property->setAccessible( true );
+		$property->setValue( $push_notifications, null );
+	}
+}
diff --git a/plugins/woocommerce/tests/php/src/Internal/PushNotifications/PushNotificationsTest.php b/plugins/woocommerce/tests/php/src/Internal/PushNotifications/PushNotificationsTest.php
index d9f1d7b489..c7e8a56085 100644
--- a/plugins/woocommerce/tests/php/src/Internal/PushNotifications/PushNotificationsTest.php
+++ b/plugins/woocommerce/tests/php/src/Internal/PushNotifications/PushNotificationsTest.php
@@ -46,6 +46,11 @@ class PushNotificationsTest extends WC_Unit_Test_Case {
 		global $wp_rest_server;
 		$wp_rest_server = null;

+		// Unregister the push token post type if it was registered.
+		if ( post_type_exists( PushToken::POST_TYPE ) ) {
+			unregister_post_type( PushToken::POST_TYPE );
+		}
+
 		$this->reset_container_replacements();
 		wc_get_container()->reset_all_resolved();

@@ -166,9 +171,22 @@ class PushNotificationsTest extends WC_Unit_Test_Case {
 	}

 	/**
-	 * @testdox Tests that register() hooks register_post_types to init when enabled.
+	 * @testdox Tests that register() hooks on_init to init action.
+	 */
+	public function test_it_hooks_on_init_to_init_action() {
+		$push_notifications = new PushNotifications();
+		$push_notifications->register();
+
+		$callback_priority = has_action( 'init', array( $push_notifications, 'on_init' ) );
+
+		$this->assertTrue( (bool) $callback_priority, 'on_init should be hooked to init' );
+		$this->assertEquals( 10, $callback_priority, 'on_init should have priority 10' );
+	}
+
+	/**
+	 * @testdox Tests that on_init registers post types when enabled.
 	 */
-	public function test_it_hooks_register_post_types_when_enabled() {
+	public function test_on_init_registers_post_types_when_enabled() {
 		$this->set_up_jetpack_connection_manager_mock( array( 'is_connected' ) );

 		$this->jetpack_connection_manager_mock
@@ -177,12 +195,9 @@ class PushNotificationsTest extends WC_Unit_Test_Case {
 			->willReturn( true );

 		$push_notifications = new PushNotifications();
-		$push_notifications->register();
-
-		$callback_priority = has_action( 'init', array( $push_notifications, 'register_post_types' ) );
+		$push_notifications->on_init();

-		$this->assertTrue( (bool) $callback_priority, 'register_post_types should be hooked to init' );
-		$this->assertEquals( 10, $callback_priority, 'register_post_types should have priority 10' );
+		$this->assertTrue( post_type_exists( PushToken::POST_TYPE ), 'Push token post type should be registered' );
 	}

 	/**
@@ -203,18 +218,17 @@ class PushNotificationsTest extends WC_Unit_Test_Case {
 	}

 	/**
-	 * @testdox Tests that push_token post type is not registered when disabled.
+	 * @testdox Tests that on_init does not register post types when disabled.
 	 */
-	public function test_it_does_not_register_push_token_post_type_when_disabled() {
+	public function test_on_init_does_not_register_post_types_when_disabled() {
 		$this->set_up_features_controller_mock( false );
-		$this->set_up_jetpack_connection_manager_mock( array( 'is_connected' ) );

 		$push_notifications = new PushNotifications();
-		$push_notifications->register();
+		$push_notifications->on_init();

 		$this->assertFalse(
-			has_action( 'init', array( $push_notifications, 'register_post_types' ) ),
-			'register_post_types should not be hooked to init when disabled'
+			post_type_exists( PushToken::POST_TYPE ),
+			'Push token post type should not be registered when disabled'
 		);
 	}