diff --git a/src/FeaturesGenerator.php b/src/FeaturesGenerator.php
index 525d781..bf199cd 100644
--- a/src/FeaturesGenerator.php
+++ b/src/FeaturesGenerator.php
@@ -153,7 +153,7 @@ class FeaturesGenerator implements FeaturesGeneratorInterface {
         }
 
         // Set the bundle machine name.
-        $packages[$package->getMachineName()]['bundle'] = $bundle->getMachineName();
+        $packages[$package->getMachineName()]->setBundle($bundle->getMachineName());
         $new_package_names[] = $package->getMachineName();
       }
       $this->featuresManager->setPackages($packages);
diff --git a/tests/src/Kernel/FeaturesGenerateTest.php b/tests/src/Kernel/FeaturesGenerateTest.php
index 6583acb..a4387d0 100644
--- a/tests/src/Kernel/FeaturesGenerateTest.php
+++ b/tests/src/Kernel/FeaturesGenerateTest.php
@@ -2,6 +2,7 @@
 
 namespace Drupal\Tests\features\Kernel;
 
+use Drupal\features\Entity\FeaturesBundle;
 use Drupal\KernelTests\KernelTestBase;
 use org\bovigo\vfs\vfsStream;
 
@@ -72,6 +73,33 @@ class FeaturesGenerateTest extends KernelTestBase {
   }
 
   /**
+   * @covers \Drupal\features\FeaturesGenerator::setPackageBundleNames
+   */
+  public function testGeneratorWithBundle() {
+
+    $filename = file_directory_temp() . '/giraffe_' . self::PACKAGE_NAME . '.tar.gz';
+    if (file_exists($filename)) {
+      unlink($filename);
+    }
+    $this->assertFalse(file_exists($filename), 'Archive file already exists.');
+
+    $bundle = FeaturesBundle::create([
+      'machine_name' => 'giraffe'
+    ]);
+
+    $this->generator->generatePackages('archive', [self::PACKAGE_NAME], $bundle);
+
+    $package = $this->featuresManager->getPackage(self::PACKAGE_NAME);
+    $this->assertNull($package);
+
+    $package = $this->featuresManager->getPackage('giraffe_' . self::PACKAGE_NAME);
+    $this->assertEquals('giraffe_' . self::PACKAGE_NAME, $package->getMachineName());
+    $this->assertEquals('giraffe', $package->getBundle());
+
+    $this->assertTrue(file_exists($filename), 'Archive file was not generated.');
+  }
+
+  /**
    * @covers \Drupal\features\Plugin\FeaturesGeneration\FeaturesGenerationWrite
    */
   public function testExportWrite() {