diff --git a/field_group.module b/field_group.module
index 5dffa87..e85c489 100644
--- a/field_group.module
+++ b/field_group.module
@@ -557,14 +557,18 @@ function field_group_get_content_element_key($context = 'default') {
 
 /**
  * Saves a group definition.
- * This function is called by ctools export when calls are made
- * through ctools_export_crud_save().
  *
- * @param $group
+ * This function is called by ctools export when calls are made through
+ * ctools_export_crud_save(). It's also used as an api method to add groups to a
+ * display.
+ *
+ * @param \stdClass $group
  *   A group definition.
- * @param $display
+ * @param \Drupal\Core\Entity\Display\EntityDisplayInterface $display
  *   The display to update if known.
- * @return EntityDisplayInterface || NULL
+ *
+ * @return \Drupal\Core\Entity\Display\EntityDisplayInterface|NULL
+ *   The updated entity display.
  */
 function field_group_group_save($group, $display = NULL) {
   if ($display === NULL) {
@@ -578,7 +582,6 @@ function field_group_group_save($group, $display = NULL) {
 
   // If no display was found. It doesn't exist yet, create it.
   if (!isset($display)) {
-
     if ($group->context == 'form') {
       $display = EntityFormDisplay::create(array(
         'targetEntityType' => $group->entity_type,
@@ -596,9 +599,6 @@ function field_group_group_save($group, $display = NULL) {
 
   }
 
-  /**
-   * @var $display \Drupal\Core\Entity\Display\EntityDisplayInterface
-   */
   if (isset($display)) {
     $data = (array) $group;
     unset($data['group_name'], $data['entity_type'], $data['bundle'], $data['mode'], $data['form'], $data['context']);
diff --git a/tests/src/Functional/EntityDisplayTest.php b/tests/src/Functional/EntityDisplayTest.php
index d83b8cd..5fc5da9 100644
--- a/tests/src/Functional/EntityDisplayTest.php
+++ b/tests/src/Functional/EntityDisplayTest.php
@@ -201,7 +201,6 @@ public function testFieldset() {
     // Test group ids and classes.
     $this->assertFieldByXPath("//fieldset[contains(@id, 'fieldset-id')]", NULL, 'Correct id set on the fieldset');
     $this->assertFieldByXPath("//fieldset[contains(@class, 'test-class')]", NULL, 'Test class set on the fieldset');
-
   }
 
   /**
@@ -278,14 +277,12 @@ public function testTabs() {
 
     // Test if it's a horizontal tab.
     $this->assertFieldByXPath('//div[@data-horizontal-tabs-panes=""]', NULL, 'Tabs are shown horizontal.');
-
   }
 
   /**
    * Test the accordion formatter.
    */
   public function testAccordion() {
-
     $data = array(
       'label' => 'Accordion item 1',
       'weight' => '1',
diff --git a/tests/src/Functional/FieldGroupTestTrait.php b/tests/src/Functional/FieldGroupTestTrait.php
index 81fa458..33683a1 100644
--- a/tests/src/Functional/FieldGroupTestTrait.php
+++ b/tests/src/Functional/FieldGroupTestTrait.php
@@ -23,9 +23,10 @@
    * @param array $data
    *   Data for the field group.
    *
-   * @return \Drupal\Core\Entity\Display\EntityDisplayInterface
+   * @return \stdClass
+   *   An object that represents the field group.
    */
-  function createGroup($entity_type, $bundle, $context, $mode, array $data) {
+  protected function createGroup($entity_type, $bundle, $context, $mode, array $data) {
 
     if (!isset($data['format_settings'])) {
       $data['format_settings'] = array();