diff --git a/tests/Database/SelectFields/UnionTests/SearchQuery.php b/tests/Database/SelectFields/UnionTests/SearchQuery.php
index aecfce53..8f04e9cc 100644
--- a/tests/Database/SelectFields/UnionTests/SearchQuery.php
+++ b/tests/Database/SelectFields/UnionTests/SearchQuery.php
@@ -41,13 +41,13 @@ public function resolve($root, $args, $ctx, ResolveInfo $info, Closure $getSelec
         if (Str::startsWith($args['id'], 'comment:')) {
             return Comment::select($selectFields->getSelect())
                 ->with($selectFields->getRelations())
-                ->where('id', (int)Str::after($args['id'], 'comment:'))
-                ->first();
-        } else {
-            return Post::select($selectFields->getSelect())
-                ->with($selectFields->getRelations())
-                ->where('id', (int)$args['id'])
+                ->where('id', (int) Str::after($args['id'], 'comment:'))
                 ->first();
         }
+
+        return Post::select($selectFields->getSelect())
+            ->with($selectFields->getRelations())
+            ->where('id', (int) $args['id'])
+            ->first();
     }
 }