diff --git a/setup/src/Magento/Setup/Test/Unit/Console/Command/UpgradeCommandTest.php b/setup/src/Magento/Setup/Test/Unit/Console/Command/UpgradeCommandTest.php
index 6209df88bd5f06880fad8c645128af079d61b482..89ea8199d6242fe299c41cb1f93d6f0adb7ec17c 100644
--- a/setup/src/Magento/Setup/Test/Unit/Console/Command/UpgradeCommandTest.php
+++ b/setup/src/Magento/Setup/Test/Unit/Console/Command/UpgradeCommandTest.php
@@ -11,7 +11,12 @@ use Magento\Framework\Console\Cli;
 
 class UpgradeCommandTest extends \PHPUnit_Framework_TestCase
 {
-    public function testExecute()
+    /**
+     * @param array $options
+     * @param string $expectedString
+     * @dataProvider executeDataProvider
+     */
+    public function testExecute($options = [], $expectedString ='')
     {
         $installerFactory = $this->getMock(\Magento\Setup\Model\InstallerFactory::class, [], [], '', false);
         $installer = $this->getMock(\Magento\Setup\Model\Installer::class, [], [], '', false);
@@ -20,6 +25,25 @@ class UpgradeCommandTest extends \PHPUnit_Framework_TestCase
         $installer->expects($this->at(2))->method('installDataFixtures');
         $installerFactory->expects($this->once())->method('create')->willReturn($installer);
         $commandTester = new CommandTester(new UpgradeCommand($installerFactory));
-        $this->assertSame(Cli::RETURN_SUCCESS, $commandTester->execute([]));
+        $this->assertSame(Cli::RETURN_SUCCESS, $commandTester->execute($options));
+        $this->assertEquals($expectedString, $commandTester->getDisplay());
+    }
+
+    /**
+     * @return array
+     */
+    public function executeDataProvider()
+    {
+        return [
+            [
+                'options' => [],
+                'expectedString' => 'Please re-run Magento compile command. Use the command "setup:di:compile"'
+                    . PHP_EOL
+            ],
+            [
+                'options' => ['--keep-generated' => true],
+                'expectedString' => ''
+            ],
+        ];
     }
 }