diff --git a/lib/behat/classes/behat_generator_base.php b/lib/behat/classes/behat_generator_base.php
index 44ea79d81c7..04942c3c965 100644
--- a/lib/behat/classes/behat_generator_base.php
+++ b/lib/behat/classes/behat_generator_base.php
@@ -25,10 +25,11 @@
 
 defined('MOODLE_INTERNAL') || die();
 
+require_once(__DIR__ . '/../../behat/behat_base.php');
+
 use Behat\Gherkin\Node\TableNode as TableNode;
 use Behat\Behat\Tester\Exception\PendingException as PendingException;
 
-
 /**
  * Class to quickly create Behat test data using component data generators.
  *
@@ -480,53 +481,7 @@ abstract class behat_generator_base {
      * @return context
      */
     protected function get_context($levelname, $contextref) {
-        global $DB;
-
-        // Getting context levels and names (we will be using the English ones as it is the test site language).
-        $contextlevels = context_helper::get_all_levels();
-        $contextnames = array();
-        foreach ($contextlevels as $level => $classname) {
-            $contextnames[context_helper::get_level_name($level)] = $level;
-        }
-
-        if (empty($contextnames[$levelname])) {
-            throw new Exception('The specified "' . $levelname . '" context level does not exist');
-        }
-        $contextlevel = $contextnames[$levelname];
-
-        // Return it, we don't need to look for other internal ids.
-        if ($contextlevel == CONTEXT_SYSTEM) {
-            return context_system::instance();
-        }
-
-        switch ($contextlevel) {
-
-            case CONTEXT_USER:
-                $instanceid = $DB->get_field('user', 'id', array('username' => $contextref));
-                break;
-
-            case CONTEXT_COURSECAT:
-                $instanceid = $DB->get_field('course_categories', 'id', array('idnumber' => $contextref));
-                break;
-
-            case CONTEXT_COURSE:
-                $instanceid = $DB->get_field('course', 'id', array('shortname' => $contextref));
-                break;
-
-            case CONTEXT_MODULE:
-                $instanceid = $DB->get_field('course_modules', 'id', array('idnumber' => $contextref));
-                break;
-
-            default:
-                break;
-        }
-
-        $contextclass = $contextlevels[$contextlevel];
-        if (!$context = $contextclass::instance($instanceid, IGNORE_MISSING)) {
-            throw new Exception('The specified "' . $contextref . '" context reference does not exist');
-        }
-
-        return $context;
+        return behat_base::get_context($levelname, $contextref);
     }
 
     /**
diff --git a/lib/behat/classes/behat_session_trait.php b/lib/behat/classes/behat_session_trait.php
index ebded0b9d61..8d2b1ad4d50 100644
--- a/lib/behat/classes/behat_session_trait.php
+++ b/lib/behat/classes/behat_session_trait.php
@@ -1068,6 +1068,69 @@ EOF;
 
         \core\session\manager::set_user($user);
     }
+
+    /**
+     * Gets the internal moodle context id from the context reference.
+     *
+     * The context reference changes depending on the context
+     * level, it can be the system, a user, a category, a course or
+     * a module.
+     *
+     * @throws Exception
+     * @param string $levelname The context level string introduced by the test writer
+     * @param string $contextref The context reference introduced by the test writer
+     * @return context
+     */
+    public static function get_context(string $levelname, string $contextref): context {
+        global $DB;
+
+        // Getting context levels and names (we will be using the English ones as it is the test site language).
+        $contextlevels = context_helper::get_all_levels();
+        $contextnames = array();
+        foreach ($contextlevels as $level => $classname) {
+            $contextnames[context_helper::get_level_name($level)] = $level;
+        }
+
+        if (empty($contextnames[$levelname])) {
+            throw new Exception('The specified "' . $levelname . '" context level does not exist');
+        }
+        $contextlevel = $contextnames[$levelname];
+
+        // Return it, we don't need to look for other internal ids.
+        if ($contextlevel == CONTEXT_SYSTEM) {
+            return context_system::instance();
+        }
+
+        switch ($contextlevel) {
+
+            case CONTEXT_USER:
+                $instanceid = $DB->get_field('user', 'id', array('username' => $contextref));
+                break;
+
+            case CONTEXT_COURSECAT:
+                $instanceid = $DB->get_field('course_categories', 'id', array('idnumber' => $contextref));
+                break;
+
+            case CONTEXT_COURSE:
+                $instanceid = $DB->get_field('course', 'id', array('shortname' => $contextref));
+                break;
+
+            case CONTEXT_MODULE:
+                $instanceid = $DB->get_field('course_modules', 'id', array('idnumber' => $contextref));
+                break;
+
+            default:
+                break;
+        }
+
+        $contextclass = $contextlevels[$contextlevel];
+        if (!$context = $contextclass::instance($instanceid, IGNORE_MISSING)) {
+            throw new Exception('The specified "' . $contextref . '" context reference does not exist');
+        }
+
+        return $context;
+    }
+
     /**
      * Trigger click on node via javascript instead of actually clicking on it via pointer.
      *
diff --git a/lib/tests/behat/behat_permissions.php b/lib/tests/behat/behat_permissions.php
index b0bca4842f6..155ad85d001 100644
--- a/lib/tests/behat/behat_permissions.php
+++ b/lib/tests/behat/behat_permissions.php
@@ -249,4 +249,44 @@ class behat_permissions extends behat_base {
             }
         }
     }
+
+    /**
+     * Mark context as frozen.
+     *
+     * @Then /^the "(?P<element_string>(?:[^"]|\\")*)" "(?P<selector_string>[^"]*)" is context frozen$/
+     * @throws ExpectationException if the context cannot be frozen or found
+     * @param string $element Element we look on
+     * @param string $selector The type of where we look (activity, course)
+     */
+    public function the_context_is_context_frozen(string $element, string $selector) {
+
+        // Enable context freeze if it is not done yet.
+        set_config('contextlocking', 1);
+
+        // Find context.
+        $context = self::get_context($selector, $element);
+
+        // Freeze context.
+        $context->set_locked(true);
+    }
+
+    /**
+     * Unmark context as frozen.
+     *
+     * @Then /^the "(?P<element_string>(?:[^"]|\\")*)" "(?P<selector_string>[^"]*)" is not context frozen$/
+     * @throws ExpectationException if the context cannot be frozen or found
+     * @param string $element Element we look on
+     * @param string $selector The type of where we look (activity, course)
+     */
+    public function the_context_is_not_context_frozen(string $element, string $selector) {
+
+        // Enable context freeze if it is not done yet.
+        set_config('contextlocking', 1);
+
+        // Find context.
+        $context = self::get_context($selector, $element);
+
+        // Freeze context.
+        $context->set_locked(false);
+    }
 }