diff --git a/src/Notifications/Notification.php b/src/Notifications/Notification.php
index 8ac5a0e3..fcaba51b 100644
--- a/src/Notifications/Notification.php
+++ b/src/Notifications/Notification.php
@@ -42,10 +42,8 @@ class Notification
 
     /**
      * Indicate that the notification gives information about a successful operation.
-     *
-     * @return $this
      */
-    public function success()
+    public function success(): self
     {
         $this->level = 'success';
 
@@ -54,10 +52,8 @@ public function success()
 
     /**
      * Indicate that the notification gives information about an error.
-     *
-     * @return $this
      */
-    public function error()
+    public function error(): self
     {
         $this->level = 'error';
 
@@ -66,18 +62,15 @@ public function error()
 
     /**
      * Set the "level" of the notification (success, error, etc.).
-     *
-     * @return $this
      */
-    public function level(string $level)
+    public function level(string $level): self
     {
         $this->level = $level;
 
         return $this;
     }
 
-    /** @return $this */
-    public function message(string $message)
+    public function message(string $message): self
     {
         $this->message = $message;
 
@@ -86,10 +79,8 @@ public function message(string $message)
 
     /**
      * Configure the "call to action" button.
-     *
-     * @return $this
      */
-    public function action(string $text, string $url)
+    public function action(string $text, string $url): self
     {
         $this->actionText = $text;
         $this->actionUrl  = $url;
@@ -97,7 +88,7 @@ public function action(string $text, string $url)
         return $this;
     }
 
-    public function to(mixed $user): Notification
+    public function to(mixed $user): self
     {
         $this->user = $user;
 
diff --git a/tests/Feature/Notifications/NotificationTest.php b/tests/Feature/Notifications/NotificationTest.php
new file mode 100644
index 00000000..b00c11bd
--- /dev/null
+++ b/tests/Feature/Notifications/NotificationTest.php
@@ -0,0 +1,55 @@
+<?php
+
+use Doctrine\ORM\EntityManagerInterface;
+use Mockery\Mock;
+use PHPUnit\Framework\TestCase;
+
+class NotificationTest extends TestCase
+{
+    /**
+     * @var Mock
+     */
+    private $registry;
+
+    /**
+     * @var Mock
+     */
+    private $em;
+
+    public function setUp(): void
+    {
+        $this->em = Mockery::spy(EntityManagerInterface::class);
+    }
+
+    public function testClassFunctions()
+    {
+        $entity = new \LaravelDoctrine\ORM\Notifications\Notification();
+
+        $entity->success();
+        $this->assertEquals('success', $entity->getLevel());
+
+        $entity->error();
+        $this->assertEquals('error', $entity->getLevel());
+
+        $entity->level('custom');
+        $this->assertEquals('custom', $entity->getLevel());
+
+        $entity->message('custom');
+        $this->assertEquals('custom', $entity->getMessage());
+
+        $entity->action('custom', 'url');
+        $this->assertEquals('custom', $entity->getActionText());
+        $this->assertEquals('url', $entity->getActionUrl());
+
+        $user = new stdClass();
+        $entity->to($user);
+        $this->assertSame($user, $entity->getUser());
+
+        $reflection = new ReflectionClass($entity);
+        $property = $reflection->getProperty('id');
+        $property->setAccessible(true);
+        $property->setValue($entity, 1);
+
+        $entity->getId();
+    }
+}