tests: include more tests

Author: luabida

poetry.lock

@@ -31,7 +31,8 @@ pydantic = "^2.12.5"
 duckdb = "^1.4.4"
 duckdb-engine = "^0.17.0"
 sqlalchemy = "^2.0.48"
-python-magic = "^0.4.27"
+python-magic = { version = "*", platform = "!=win32" }
+python-magic-bin = { version = "*", platform = "win32" }
 chardet = "^7.4.0.post2"
 anyio = "^4.13.0"
 httpx = ">=0.28.0"

pyproject.toml

@@ -26,7 +26,7 @@
 from .ftp import FTPClient
 from .models import BaseLocalFile, BaseRemoteFile
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from duckdb import DuckDBPyConnection
 
 
@@ -403,7 +403,7 @@ async def download_to_parquet(
 
         if hasattr(local_file, "to_parquet"):
             original_path = local_file.path
-            parquet_file = await local_file.to_parquet(callback=callback)  # type: ignore
+            parquet_file = await local_file.to_parquet(callback=callback)
             parquet_file.add_dv = add_dv
 
             await self._update_state(

pysus/api/client.py

@@ -13,7 +13,7 @@
 from pysus.api.models import BaseRemoteClient, BaseRemoteFile
 from pysus.api.types import DADOSGOV
 
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from .models import Dataset
 
 

pysus/api/dadosgov/client.py

@@ -272,7 +272,7 @@ def formatter(self, filename: str) -> dict[str, Any]:
                 }
 
             m = re.search(r"_(\w{3})-out_(\d{4})_\.csv$", name)
-            if m:
+            if m:  # pragma: no cover
                 return {
                     "state": None,
                     "year": _parse_year(m.group(2)),

pysus/api/dadosgov/databases.py

@@ -74,12 +74,12 @@ class Dataset(Base):
     description = Column(String, nullable=True)
 
     groups = relationship(
-        "DatasetGroup",
+        "Group",
         back_populates="dataset",
         cascade="all, delete-orphan",
     )
     files = relationship(
-        "CatalogFile",
+        "File",
         back_populates="dataset",
         cascade="all, delete-orphan",
     )
@@ -128,9 +128,9 @@ class ColumnDefinition(Base):
     description = Column(String, nullable=True)
     nullable = Column(Boolean, nullable=False, default=True)
 
-    dataset = relationship("CatalogDataset", back_populates="columns")
+    dataset = relationship("Dataset", back_populates="columns")
     files = relationship(
-        "CatalogFile",
+        "File",
         secondary=file_columns,
         back_populates="columns",
     )
@@ -177,11 +177,11 @@ class Group(Base):
     description = Column(String, nullable=True)
 
     dataset = relationship(
-        "CatalogDataset",
+        "Dataset",
         back_populates="groups",
     )
     files = relationship(
-        "CatalogFile",
+        "File",
         back_populates="group",
         cascade="all, delete-orphan",
     )
@@ -281,11 +281,11 @@ class File(Base):
     )
 
     dataset: Mapped["Dataset"] = relationship(
-        "CatalogDataset",
+        "Dataset",
         back_populates="files",
     )
     group: Mapped[Optional["Group"]] = relationship(
-        "DatasetGroup",
+        "Group",
         back_populates="files",
     )
     columns: Mapped[list["ColumnDefinition"]] = relationship(

pysus/api/ducklake/catalog/orm/dataset.py

@@ -15,7 +15,8 @@ class Base(DeclarativeBase):
 
 
 class Dataset(Base):
-    """ORM model for the datasets table — central registry of available datasets.
+    """ORM model for the datasets table — central registry of available
+    datasets.
 
     Parameters
     ----------

pysus/api/ducklake/catalog/orm/default.py

@@ -77,8 +77,8 @@ def __init__(self, engine=None, **data) -> None:
         self._engine = engine
         self._cache_dir: Path = Path(CACHEPATH) / "ducklake"
         self._cache_dir.mkdir(parents=True, exist_ok=True)
-        self._catalog_local: Path = self._cache_dir / "catalog.db"
-        self._catalog_remote: str = "public/catalog.db"
+        self._catalog_local: Path = self._cache_dir / "catalog.duckdb"
+        self._catalog_remote: str = "public/catalog.duckdb"
 
     @property
     def name(self) -> str:
@@ -197,7 +197,8 @@ def _setup_engine(self, local_path: Path | None = None):
         Parameters
         ----------
         local_path : Path, optional
-            Path to the catalog database file. Defaults to the discovery catalog.
+            Path to the catalog database file.
+            Defaults to the discovery catalog.
         """
         if local_path is None:
             local_path = self._catalog_local
@@ -210,7 +211,8 @@ def _setup_engine(self, local_path: Path | None = None):
             conn.exec_driver_sql("INSTALL ducklake; LOAD ducklake;")
 
             has_pysus = conn.exec_driver_sql(
-                "SELECT 1 FROM information_schema.schemata WHERE schema_name = 'pysus'"
+                "SELECT 1 FROM information_schema.schemata"
+                " WHERE schema_name = 'pysus'"
             ).fetchone()
 
             if has_pysus:
@@ -326,7 +328,9 @@ async def _download(
                 else:
                     raise e
 
-    async def _download_catalog(self, local_path: Path, remote_path: str) -> None:
+    async def _download_catalog(
+        self, local_path: Path, remote_path: str
+    ) -> None:
         """Download a catalog database from remote storage with retries.
 
         Parameters
@@ -380,7 +384,9 @@ def _get_s3_client(self):
             "s3",
             endpoint_url=f"https://{self.endpoint}",
             aws_access_key_id=self.credentials.access_key.get_secret_value(),
-            aws_secret_access_key=(self.credentials.secret_key.get_secret_value()),
+            aws_secret_access_key=(
+                self.credentials.secret_key.get_secret_value()
+            ),
             region_name=self.region,
             config=Config(signature_version="s3v4"),
         )
@@ -400,7 +406,10 @@ async def _upload_catalog(self) -> None:
             if not ds._catalog_local.exists():
                 continue
 
-            def _upload(local=str(ds._catalog_local), name=ds._catalog_name):
+            _local = str(ds._catalog_local)
+            _name = ds._catalog_name
+
+            def _upload(local=_local, name=_name):
                 self._s3_client.upload_file(
                     local,
                     self.bucket,

pysus/api/ducklake/client.py

@@ -13,16 +13,13 @@
 from anyio import to_thread
 from pydantic import Field, PrivateAttr
 from pysus import CACHEPATH
+from pysus.api.ducklake.catalog.orm.dataset import Dataset
+from pysus.api.ducklake.catalog.orm.dataset import File as CatalogFile
+from pysus.api.ducklake.catalog.orm.dataset import Group
 from pysus.api.models import BaseRemoteDataset, BaseRemoteFile, BaseRemoteGroup
 from sqlalchemy.orm import contains_eager, joinedload, sessionmaker
 
-from pysus.api.ducklake.catalog.orm.dataset import (
-    Dataset,
-    File as CatalogFile,
-    Group,
-)
-
-if TYPE_CHECKING:
+if TYPE_CHECKING:  # pragma: no cover
     from .client import DuckLake
 
 
@@ -41,23 +38,19 @@ class File(BaseRemoteFile):
         The parent group object, if any.
     """
 
-    def __init__(
-        self,
-        dataset: "DuckDataset",
-        record: CatalogFile,
-        group: Optional["DuckGroup"] = None,
-    ) -> None:
+    record: CatalogFile = Field(exclude=True)
+    group: Optional["DuckGroup"] = Field(default=None, exclude=True)
+
+    def __init__(self, **data: Any) -> None:
+        record = data.pop("record")
+        group = data.pop("group", None)
         super().__init__(
             path=Path(record.path),
             type=record.type or "remote",
-            dataset=dataset,
+            record=record,  # type: ignore[call-arg]
+            group=group,
+            **data,
         )
-        self.record: CatalogFile = record
-        self.group: Optional["DuckGroup"] = group
-
-    @property
-    def path(self) -> Path:
-        return Path(self.record.path)
 
     @property
     def basename(self) -> str:
@@ -191,7 +184,7 @@ def __init__(self, **data) -> None:
         super().__init__(**data)
         self._cache_dir: Path = Path(CACHEPATH) / "ducklake"
         self._cache_dir.mkdir(parents=True, exist_ok=True)
-        self._catalog_name: str = f"catalog_{self.record.name.lower()}.db"
+        self._catalog_name: str = f"catalog_{self.record.name.lower()}.duckdb"
         self._catalog_local: Path = self._cache_dir / self._catalog_name
 
     def __repr__(self) -> str:
@@ -306,7 +299,7 @@ def _upload():
             self.client._s3_client.upload_file(
                 str(self._catalog_local),
                 self.client.bucket,
-                f"catalog_{self.record.name.lower()}.db",
+                f"catalog_{self.record.name.lower()}.duckdb",
             )
 
         await to_thread.run_sync(_upload)