diff --git a/tests/routers/openml/datasets_test.py b/tests/routers/openml/datasets_test.py
index 1a5897b..53e70e4 100644
--- a/tests/routers/openml/datasets_test.py
+++ b/tests/routers/openml/datasets_test.py
@@ -175,6 +175,24 @@ async def test_dataset_features(py_api: httpx.AsyncClient) -> None:
     ]
 
 
+async def test_dataset_features_with_ontology(py_api: httpx.AsyncClient) -> None:
+    # Dataset 11 has ontology data for features 1, 2, and 3
+    response = await py_api.get("/datasets/features/11")
+    assert response.status_code == HTTPStatus.OK
+    features = {f["index"]: f for f in response.json()}
+    assert features[1]["ontology"] == ["https://en.wikipedia.org/wiki/Service_(motor_vehicle)"]
+    assert features[2]["ontology"] == [
+        "https://en.wikipedia.org/wiki/Car_door",
+        "https://en.wikipedia.org/wiki/Door",
+    ]
+    assert features[3]["ontology"] == [
+        "https://en.wikipedia.org/wiki/Passenger_vehicles_in_the_United_States"
+    ]
+    # Features without ontology should not include the field
+    assert "ontology" not in features[0]
+    assert "ontology" not in features[4]
+
+
 async def test_dataset_features_no_access(py_api: httpx.AsyncClient) -> None:
     response = await py_api.get("/datasets/features/130")
     assert response.status_code == HTTPStatus.FORBIDDEN
diff --git a/tests/routers/openml/migration/datasets_migration_test.py b/tests/routers/openml/migration/datasets_migration_test.py
index 75f3086..2abaad1 100644
--- a/tests/routers/openml/migration/datasets_migration_test.py
+++ b/tests/routers/openml/migration/datasets_migration_test.py
@@ -260,11 +260,11 @@ async def test_datasets_feature_is_identical(
                 # The old API returns a str if there is only a single element
                 feature["nominal_value"] = values if len(values) > 1 else values[0]
             elif key == "ontology":
-                del feature[key]  # Added back in with follow up PR #262
+                # The old API returns a str if there is only a single element
+                values = feature.pop(key)
+                feature["ontology"] = values if len(values) > 1 else values[0]
             else:
                 # The old API formats bool as string in lower-case
                 feature[key] = str(value) if not isinstance(value, bool) else str(value).lower()
     original_features = original.json()["data_features"]["feature"]
-    for feature in original_features:
-        feature.pop("ontology", None)
     assert python_body == original_features