add tests for ndjson

Author: vincentsarago

tests/routes/test_items.py

@@ -561,9 +561,7 @@ def test_output_response_type(app):
     body = response.json()
     assert len(body) == 10
     feat = body[0]
-    assert ["collectionId", "itemId", "id", "pr", "row", "path", "ogc_fid"] == list(
-        feat.keys()
-    )
+    assert "geometry" not in feat.keys()
 
     response = app.get(
         "/collections/public.landsat_wrs/items",
@@ -573,3 +571,30 @@ def test_output_response_type(app):
     assert response.headers["content-type"] == "application/json"
     body = response.json()
     assert len(body) == 10
+
+    # ndjson output
+    response = app.get("/collections/public.landsat_wrs/items?f=ndjson")
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "application/ndjson"
+    body = response.text.splitlines()
+    assert len(body) == 10
+    feat = json.loads(body[0])
+    assert [
+        "collectionId",
+        "itemId",
+        "id",
+        "pr",
+        "row",
+        "path",
+        "ogc_fid",
+        "geometry",
+    ] == list(feat.keys())
+
+    response = app.get(
+        "/collections/public.landsat_wrs/items",
+        headers={"accept": "application/ndjson"},
+    )
+    assert response.status_code == 200
+    assert response.headers["content-type"] == "application/ndjson"
+    body = response.text.splitlines()
+    assert len(body) == 10

tifeatures/factory.py

@@ -3,7 +3,7 @@
 import csv
 import json
 from dataclasses import dataclass, field
-from typing import Any, Callable, Dict, Iterable, List, Optional
+from typing import Any, Callable, Dict, Generator, Iterable, List, Optional
 
 import jinja2
 from pygeofilter.ast import AstType
@@ -39,30 +39,6 @@
 from starlette.responses import StreamingResponse
 from starlette.templating import Jinja2Templates, _TemplateResponse
 
-
-class DummyWriter:
-    """Dummy writer that implements write for use with csv.writer."""
-
-    def write(self, line: str):
-        """Return line."""
-        return line
-
-
-def iter_csv(data: Iterable[Dict]):
-    """Creates an iterator that returns lines of csv from an iterable of dicts."""
-
-    initial = True
-    writer = None
-    for row in data:
-        if initial:
-            fieldnames = row.keys()
-            writer = csv.DictWriter(DummyWriter(), fieldnames=fieldnames)
-            yield writer.writeheader()
-            initial = False
-        if writer:
-            yield writer.writerow(row)
-
-
 settings = APISettings()
 
 # custom template directory
@@ -112,6 +88,29 @@ def create_html_response(
     )
 
 
+def create_csv_rows(data: Iterable[Dict]) -> Generator[str, None, None]:
+    """Create Template response."""
+
+    class DummyWriter:
+        """Dummy writer that implements write for use with csv.writer."""
+
+        def write(self, line: str):
+            """Return line."""
+            return line
+
+    """Creates an iterator that returns lines of csv from an iterable of dicts."""
+    initial = True
+    writer = None
+    for row in data:
+        if initial:
+            fieldnames = row.keys()
+            writer = csv.DictWriter(DummyWriter(), fieldnames=fieldnames)
+            yield writer.writeheader()
+            initial = False
+        if writer:
+            yield writer.writerow(row)
+
+
 @dataclass
 class Endpoints:
     """Endpoints Factory."""
@@ -534,6 +533,7 @@ def queryables(
                         MediaType.csv.value: {},
                         MediaType.json.value: {},
                         MediaType.geojsonseq.value: {},
+                        MediaType.ndjson.value: {},
                     }
                 },
             },
@@ -621,7 +621,7 @@ async def items(
                 MediaType.json,
                 MediaType.ndjson,
             ):
-                if items[0].geometry is not None:
+                if items and items[0].geometry is not None:
                     rows = (
                         {
                             "collectionId": collection.id,
@@ -631,6 +631,7 @@ async def items(
                         }
                         for f in items
                     )
+
                 else:
                     rows = (
                         {
@@ -644,7 +645,7 @@ async def items(
                 # CSV Response
                 if output_type == MediaType.csv:
                     return StreamingResponse(
-                        iter_csv(rows),
+                        create_csv_rows(rows),
                         media_type=MediaType.csv,
                         headers={
                             "Content-Disposition": "attachment;filename=items.csv"
@@ -658,7 +659,7 @@ async def items(
                 # NDJSON Response
                 if output_type == MediaType.ndjson:
                     return StreamingResponse(
-                        (row + "\n" for row in rows),
+                        (json.dumps(row) + "\n" for row in rows),
                         media_type=MediaType.ndjson,
                         headers={
                             "Content-Disposition": "attachment;filename=items.ndjson"