Add or Remove extra roles from wp account upon subscription activation or cancellation

This code snippet adds and removes extra role subscriber from wp account upon subscription activation or cancellation

add_action( 'umm_stripe_status_active', 'um_03222024_stripe_swap_roles', 999999, 2 );
add_action( 'umm_stripe_status_canceled', 'um_03222024_stripe_swap_roles', 999999, 2 );
function um_03222024_stripe_swap_roles( $um_user_id, $role_active ){
 
    $stripe        = UM()->Stripe_API()->common()->stripe_api();
    $customer_id   = UM()->Stripe_API()->common()->request_handler()->customer_id;
    $subscriptions = $stripe->subscriptions->all(['limit' => 3, 'customer' => $customer_id ]);

    $extra_role = 'subscriber';

    // Remove 'subscriber' role when active subscriptions is more than 1
    if( isset( $subscriptions->data ) && count( $subscriptions->data ) > 1 ) {
        UM()->roles()->remove_role( $um_user_id, sanitize_key( $extra_role ) );
    }

    //Add 'subscriber role when active subscriptions is less than or equal to 1
    if( isset( $subscriptions->data ) && count( $subscriptions->data ) <= 1 ) {
        $user = get_user_by( 'ID', $um_user_id );
        $user->add_role( $extra_role );
    }
}