test: test _extract content from docset

tests/unit/types/sets/test_documentset.py

@@ -1,6 +1,7 @@
 from copy import deepcopy
 
 import pytest
+import numpy as np
 
 from jina import Document
 from jina.types.sets import DocumentSet
@@ -220,3 +221,21 @@ def callback_fn(docs, *args, **kwargs) -> None:
         add_chunk(doc)
         add_match(doc)
         add_match(doc)
+
+
+@pytest.mark.parametrize('num_rows', [1, 2, 3])
+@pytest.mark.parametrize('field', ['content', 'blob', 'embedding'])
+def test_get_content(num_rows, field):
+    batch_size = 10
+    embed_size = 20
+
+    kwargs = {
+        field: np.random.random((num_rows, embed_size))
+    }
+
+    docs = DocumentSet([Document(**kwargs) for _ in range(batch_size)])
+    docs.append(Document())
+
+    contents, pts = docs._extract_docs(field)
+
+    assert contents.shape == (batch_size, num_rows, embed_size)