Shipping rate test plugin for checkout blocks bug(?)

shipping-cache-test-block.php

<?php

/*
Plugin Name: shipping-cache-test-block
Version: 1.0
*/

if ( ! defined( 'ABSPATH' ) ) {
	exit;
}

//Activate BACS and COD payment method. Switching between these two will trigger the shipping rate to change. 
add_action('woocommerce_shipping_init', function(){
    class Shipping_Cache_Test_Block extends WC_Shipping_Method {

        public function __construct( $instance_id = 0 ) {
            $this->id                 = 'test_shipping_rate';
            $this->instance_id        = absint( $instance_id );
            $this->method_title       = 'test rate';
            $this->supports           = array('shipping-zones');
            $this->init();

        }

        private function init() {
            $this->title = 'test rate';
        }

        public function calculate_shipping( $package = array() ) {
            $rate = array(
                'id'      => $this->get_rate_id(),
                'label'   => $this->title,
                'cost'    => 100,
                'package' => $package,
            );

            $payment_method = WC()->session->get('custom_chosen_payment_method');
            if($payment_method == 'cod') {
                $rate['cost'] = 200;
            }

            $this->add_rate( $rate );
        }

    }

    add_filter( 'woocommerce_shipping_methods', function(){
        $methods['test_shipping_rate'] = 'Shipping_Cache_Test_Block';
        return $methods;
    });
});

add_action('woocommerce_blocks_loaded', function(){
    woocommerce_store_api_register_update_callback([
        'namespace' => 'test_shipping_rate',
        'callback'  => function( $data ) {
            if(isset($data['payment_method'])) {
                WC()->session->set('custom_chosen_payment_method', $data['payment_method']);
            }
        }
    ]);
});

//This is required since payment method selection is not persistent
add_action('woocommerce_checkout_init', function(){
    if ( is_admin() && ! defined( 'DOING_AJAX' ) )
    return;

    if( ( defined( 'REST_REQUEST' ) && REST_REQUEST ) || ( defined( 'DOING_AJAX') && DOING_AJAX ))
    return;

    $available_gateways = WC()->payment_gateways->get_available_payment_gateways();
    WC()->session->set('custom_chosen_payment_method', key($available_gateways));
});

add_action('wp_footer', function(){
    ?>
    <script>
        wp.hooks.addAction('experimental__woocommerce_blocks-checkout-set-active-payment-method', 'vp-shipping-rate', function(payment_method) {
            window.wc.blocksCheckout.extensionCartUpdate({
                namespace: 'test_shipping_rate',
                data: {
                    payment_method: payment_method.value
                }
            });
        });
    </script>
    <?php
});