diff --git a/src/MediaController.php b/src/MediaController.php
index cbad49c..ab5a685 100644
--- a/src/MediaController.php
+++ b/src/MediaController.php
@@ -8,6 +8,7 @@ use Illuminate\Http\Request;
 use Illuminate\Support\Facades\Storage;
 use Illuminate\Validation\Rule;
 use Illuminate\Validation\ValidationException;
+use Spatie\LaravelData\DataCollection;
 use Spatie\MediaLibrary\HasMedia;
 use Spatie\MediaLibrary\MediaCollections\MediaCollection;
 use Spatie\MediaLibrary\MediaCollections\Models\Media;
@@ -71,15 +72,15 @@ class MediaController
         return MediaData::from($media);
     }
 
-    public function index(Request $request, $parentModel, int $parentId, string $collection): JsonResponse
+    public function index(Request $request, $parentModel, int $parentId, string $collection): MediaData|DataCollection
     {
         $model = app('media-library-helpers')->get($parentModel);
         $model = $model::find($parentId);
-        $isSingle = $model->getMediaCollection($collection)->collectionSizeLimit;
+        $isSingle = 1 === $model->getMediaCollection($collection)->collectionSizeLimit;
 
-        return response()->json([
-            'data' => $isSingle ? $model->getFirstMedia($collection) : $model->getMedia($collection)->map(fn ($c) => $c->toArray()),
-        ]);
+        return $isSingle
+            ? MediaData::from($model->getFirstMedia($collection))
+            : MediaData::collection($model->getMedia($collection));
     }
 
     public function destroy(Media $media, Request $request): JsonResponse
diff --git a/tests/Feature/IndexTest.php b/tests/Feature/IndexTest.php
new file mode 100644
index 0000000..857b161
--- /dev/null
+++ b/tests/Feature/IndexTest.php
@@ -0,0 +1,29 @@
+<?php
+
+namespace Zoomyboy\MedialibraryHelper\Tests\Feature;
+
+test('it gets all medias', function () {
+    $this->auth()->registerModel();
+    $post = $this->newPost();
+    $firstMedia = $post->addMedia($this->pdfFile()->getPathname())->withCustomProperties(['test' => 'old'])->preservingOriginal()->toMediaCollection('images');
+    $secondMedia = $post->addMedia($this->pdfFile()->getPathname())->withCustomProperties(['test' => 'old'])->preservingOriginal()->toMediaCollection('images');
+
+    $response = $this->getJson("/mediaupload/post/{$post->id}/images");
+
+    $response->assertStatus(200);
+    $response->assertJsonPath('0.id', $firstMedia->id);
+    $response->assertJsonPath('1.id', $secondMedia->id);
+    $response->assertJsonPath('1.properties.test', 'old');
+});
+
+test('it gets media for single', function () {
+    $this->auth()->registerModel();
+    $post = $this->newPost();
+    $media = $post->addMedia($this->pdfFile()->getPathname())->withCustomProperties(['test' => 'old'])->preservingOriginal()->toMediaCollection('defaultSingleFile');
+
+    $response = $this->getJson("/mediaupload/post/{$post->id}/defaultSingleFile");
+
+    $response->assertStatus(200);
+    $response->assertJsonPath('id', $media->id);
+    $response->assertJsonPath('properties.test', 'old');
+});