diff --git a/mod/assign/tests/downloader_test.php b/mod/assign/tests/downloader_test.php
index d7716ec5113..da077a3a9ba 100644
--- a/mod/assign/tests/downloader_test.php
+++ b/mod/assign/tests/downloader_test.php
@@ -116,7 +116,7 @@ class downloader_test extends \advanced_testcase {
             }
             $datagenerator->create_submission([
                 'userid' => $user->id,
-                'assignid' => $cm->id,
+                'cmid' => $cm->id,
                 'file' => implode(',', $files),
             ]);
         }
diff --git a/mod/assign/tests/generator/behat_mod_assign_generator.php b/mod/assign/tests/generator/behat_mod_assign_generator.php
index 615fd04c342..8e3fd587e70 100644
--- a/mod/assign/tests/generator/behat_mod_assign_generator.php
+++ b/mod/assign/tests/generator/behat_mod_assign_generator.php
@@ -37,16 +37,22 @@ class behat_mod_assign_generator extends behat_generator_base {
                 'required' => ['assign', 'user'],
                 'switchids' => ['assign' => 'assignid', 'user' => 'userid'],
             ],
+            'extensions' => [
+                'singular' => 'extension',
+                'datagenerator' => 'extension',
+                'required' => ['assign', 'user', 'extensionduedate'],
+                'switchids' => ['assign' => 'cmid', 'user' => 'userid'],
+            ],
         ];
     }
 
     /**
-     * Get the assignment CMID using an activity idnumber.
+     * Get the assignment cmid using an activity name or idnumber.
      *
-     * @param string $idnumber
+     * @param string $identifier activity name or idnumber
      * @return int The cmid
      */
-    protected function get_assign_id(string $idnumber): int {
-        return $this->get_activity_id($idnumber);
+    protected function get_assign_id(string $identifier): int {
+        return $this->get_cm_by_activity_name('assign', $identifier)->id;
     }
 }
diff --git a/mod/assign/tests/generator/lib.php b/mod/assign/tests/generator/lib.php
index 85fe6620afd..7d834979c58 100644
--- a/mod/assign/tests/generator/lib.php
+++ b/mod/assign/tests/generator/lib.php
@@ -77,11 +77,22 @@ class mod_assign_generator extends testing_module_generator {
     /**
      * Create an assignment submission.
      *
-     * @param array $data
+     * @param array $data with keys userid, cmid and
+     *      then data for each assignsubmission plugin used.
+     *      For backwards compatibility, you can pass cmid as 'assignid' but that generates a warning.
      */
     public function create_submission(array $data): void {
         global $USER;
 
+        if (array_key_exists('assignid', $data)) {
+            debugging(
+                'The cmid passed to create_submission should have array key cmid, not assignid.',
+                DEBUG_DEVELOPER,
+            );
+            $data['cmid'] = $data['assignid'];
+            unset($data['assignid']);
+        }
+
         $currentuser = $USER;
         $user = \core_user::get_user($data['userid']);
         $this->set_user($user);
@@ -90,7 +101,7 @@ class mod_assign_generator extends testing_module_generator {
             'userid' => $user->id,
         ];
 
-        [$course, $cm] = get_course_and_cm_from_cmid($data['assignid'], 'assign');
+        [$course, $cm] = get_course_and_cm_from_cmid($data['cmid'], 'assign');
         $context = context_module::instance($cm->id);
         $assign = new assign($context, $cm, $course);
 
@@ -102,11 +113,28 @@ class mod_assign_generator extends testing_module_generator {
             }
         }
 
-        $assign->save_submission((object) $submission, $notices);
+        $assign->save_submission($submission, $notices);
 
         $this->set_user($currentuser);
     }
 
+    /**
+     * Create an assignment extension.
+     *
+     * @param array $data must have keys cmid, userid, extensionduedate.
+     */
+    public function create_extension(array $data): void {
+        $user = \core_user::get_user($data['userid'], '*', MUST_EXIST);
+
+        [$course, $cm] = get_course_and_cm_from_cmid($data['cmid'], 'assign');
+        $context = context_module::instance($cm->id);
+        $assign = new assign($context, $cm, $course);
+
+        if (!$assign->save_user_extension($user->id, $data['extensionduedate'] ?: null)) {
+            throw new \core\exception\coding_exception('The requested extension could not be created.');
+        }
+    }
+
     /**
      * Gets the grouping id from it's idnumber.
      *
@@ -119,7 +147,7 @@ class mod_assign_generator extends testing_module_generator {
 
         // Do not fetch grouping ID for empty grouping idnumber.
         if (empty($idnumber)) {
-            return null;
+            throw new \core\exception\coding_exception('idnumber cannot be empty');
         }
 
         if (!$id = $DB->get_field('groupings', 'id', ['idnumber' => $idnumber])) {