1 files changed, 38 insertions, 0 deletions
diff --git a/tests/ControllerTest.php b/tests/ControllerTest.php
index 8cdbf7d..5db4117 100644
--- a/tests/ControllerTest.php
+++ b/tests/ControllerTest.php
@@ -1,6 +1,7 @@
 <?php
 
 use Jasny\Controller;
+use Jasny\Flash;
 use Psr\Http\Message\ServerRequestInterface;
 use Psr\Http\Message\ResponseInterface;
 
@@ -76,6 +77,43 @@ class ControllerTest extends PHPUnit_Framework_TestCase
     }
 
     /**
+     * Test setting flash
+     *
+     * @dataProvider flashProvider
+     * @param object $data 
+     */
+    public function testFlash($data)
+    {
+        $controller = $this->getMockBuilder(Controller::class)->disableOriginalConstructor()->getMockForAbstractClass();
+
+        $flash = $controller->flash();
+        $this->assertInstanceOf(Flash::class, $flash, "Flash is not set");
+        $this->assertEmpty($flash->get(), "Flash data should be empty");        
+
+        $flash = $controller->flash($data->type, $data->message);
+        $this->assertInstanceOf(Flash::class, $flash, "Flash is not set");
+        $this->assertEquals($data, $flash->get(), "Flash data is incorrect");
+
+        $flash = $controller->flash();
+        $this->assertInstanceOf(Flash::class, $flash, "Flash is not set");
+        $this->assertEquals($data, $flash->get(), "Flash data is incorrect");
+
+        $flash->clear();
+    }
+
+    /**
+     * Test setting flash
+     *
+     * @return array
+     */
+    public function flashProvider()
+    {
+        return [
+            [(object)['type' => 'test_type', 'message' => 'Test message']]
+        ];  
+    }
+
+    /**
      * Get map of status codes to states
      *
      * @param int $code