Accessing other contexts in Behat 3

behat.yml

default:
  suites:
    default:
      contexts: [FeatureContext, SubContext]

FeatureContext.php

<?php

use Behat\Behat\Context\Context;
use Behat\Behat\Hook\Scope\BeforeScenarioScope;

class FeatureContext implements Context {
    private $subContext;

    /** @BeforeScenario */
    public function gatherContexts(BeforeScenarioScope $scope)
    {
        $environment = $scope->getEnvironment();
    
        $this->subContext = $environment->getContext('SubContext');
    }

    /** @When I set the state :state */
    function IGetTheContext($state) {
        $this->subContext->setState($state);
    }
}

subcontext.feature

Feature: SubContext
  Scenario: Access the sub context
    When I set the state "foo"
    Then the state should be "foo"

SubContext.php

<?php

use Behat\Behat\Context\Context;

class SubContext implements Context {
    private $state;

    function setState($state) {
        $this->state = $state;
    }

    /** @Then the state should be :state */
    function theStateShouldBe($state) {
        PHPUnit_Framework_Assert::assertSame($state, $this->state);
    }
}

Thanks!