Add tests

model2vec/utils.py

@@ -182,7 +182,7 @@ def _get_metadata_from_readme(readme_path: Path) -> dict[str, Any]:
         logger.info(f"README file not found in {readme_path}. No model card loaded.")
         return {}
     model_card = ModelCard.load(readme_path)
-    data: dict[str, Any] = model_card.data
+    data: dict[str, Any] = model_card.data.to_dict()
     if not data:
         logger.info("File README.md exists, but was empty. No model card loaded.")
     return data

tests/test_utils.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+
+from model2vec.utils import _get_metadata_from_readme
+
+
+def test__get_metadata_from_readme_not_exists() -> None:
+    """Test getting metadata from a README."""
+    assert _get_metadata_from_readme(Path("zzz")) == {}
+
+
+def test__get_metadata_from_readme_mocked_file() -> None:
+    """Test getting metadata from a README."""
+    with NamedTemporaryFile() as f:
+        f.write(b"---\nkey: value\n---\n")
+        f.flush()
+        assert _get_metadata_from_readme(Path(f.name))["key"] == "value"
+
+
+def test__get_metadata_from_readme_mocked_file_keys() -> None:
+    """Test getting metadata from a README."""
+    with NamedTemporaryFile() as f:
+        f.write(b"")
+        f.flush()
+        assert set(_get_metadata_from_readme(Path(f.name))) == set()