tests/test_fisheries_hf.py

import csv
import sys
import tempfile
import unittest
import zipfile
from pathlib import Path
from unittest.mock import MagicMock

sys.modules.setdefault("httpx", MagicMock())
import fisheries_hf


CSV_TEXT = """year,month,lon,lat,effort,flag
2022,12,141,2,9,A
2023,1,141,2,1.5,A
2023,2,142,3,2.5,B
2023,2,142,3,2.5,B
2023,3,150,3,8,C
"""


class FisheriesHFTests(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.root = Path(self.tmp.name)
        fisheries_hf.EXPORT_ROOT = self.root / "exports"
        fisheries_hf.EXPORT_ROOT.mkdir()

    def tearDown(self):
        self.tmp.cleanup()

    def _check(self, path: Path):
        result = fisheries_hf.analyze_and_export(
            path,
            dataset_path="current/test" + path.suffix,
            revision="a" * 40,
            repository="tuna",
            year=2023,
            lon_min=140,
            lon_max=145,
            lat_min=0,
            lat_max=5,
        )
        self.assertEqual(result["status"], "ok")
        self.assertEqual(
            result["repository"],
            "globalsquiddatabase/Tuna-Fisheries-Dataset",
        )
        self.assertEqual(result["scanned_row_count"], 5)
        self.assertEqual(result["matched_row_count"], 3)
        self.assertEqual(result["time_range"], {"min_year": 2023, "max_year": 2023})
        self.assertEqual(result["exact_duplicate_count"], 1)
        self.assertEqual(result["missing_values_by_column"]["effort"], 0)
        self.assertEqual(result["monthly_record_counts"], [
            {"month": 1, "record_count": 1},
            {"month": 2, "record_count": 2},
        ])
        self.assertEqual(result["annual_summary"][0]["record_count"], 3)
        self.assertAlmostEqual(result["annual_summary"][0]["sum_effort"], 6.5)
        self.assertEqual(result["deduplicated_record_count"], 2)
        self.assertEqual(
            [item["kind"] for item in result["exports"]],
            ["filtered_raw", "deduplicated", "annual_summary"],
        )
        self.assertEqual(len(result["download_urls"]), 3)
        export = fisheries_hf.EXPORT_ROOT / result["token"] / result["filename"]
        self.assertTrue(export.is_file())
        with export.open(encoding="utf-8-sig", newline="") as stream:
            rows = list(csv.DictReader(stream))
        self.assertEqual(len(rows), 3)
        for item in result["exports"]:
            exported = fisheries_hf.EXPORT_ROOT / item["token"] / item["filename"]
            self.assertTrue(exported.is_file())

    def test_csv_filter_quality_and_export(self):
        path = self.root / "sample.csv"
        path.write_text(CSV_TEXT, encoding="utf-8")
        self._check(path)

    def test_zip_csv_member_filter_and_export(self):
        path = self.root / "sample.zip"
        with zipfile.ZipFile(path, "w", compression=zipfile.ZIP_DEFLATED) as archive:
            archive.writestr("monthly_2023.csv", CSV_TEXT)
        self._check(path)

    def test_rejects_invalid_bbox(self):
        path = self.root / "sample.csv"
        path.write_text(CSV_TEXT, encoding="utf-8")
        with self.assertRaisesRegex(ValueError, "lon_min"):
            fisheries_hf.analyze_and_export(
                path,
                dataset_path="current/test.csv",
                revision="a" * 40,
                lon_min=181,
            )

    def test_repository_allowlist(self):
        self.assertEqual(
            fisheries_hf.normalize_repository("squid"),
            "globalsquiddatabase/squid_dataset",
        )
        self.assertEqual(
            fisheries_hf.normalize_repository("tuna"),
            "globalsquiddatabase/Tuna-Fisheries-Dataset",
        )
        with self.assertRaisesRegex(ValueError, "repository"):
            fisheries_hf.normalize_repository("someone/other-dataset")

    def test_local_mirror_overrides_only_matching_remote_path(self):
        local_root = self.root / "local-squid"
        local_root.mkdir()
        (local_root / "cached.csv").write_text("year\n2023\n", encoding="utf-8")
        previous = fisheries_hf.LOCAL_DATA_ROOTS["squid"]
        fisheries_hf.LOCAL_DATA_ROOTS["squid"] = local_root
        try:
            merged = fisheries_hf.merge_local_with_remote_tree(
                [
                    {"path": "cached.csv", "type": "file", "size": 999},
                    {"path": "remote-only.csv", "type": "file", "size": 42},
                ],
                "squid",
            )
        finally:
            fisheries_hf.LOCAL_DATA_ROOTS["squid"] = previous
        by_path = {item["path"]: item for item in merged}
        self.assertEqual(set(by_path), {"cached.csv", "remote-only.csv"})
        self.assertEqual(by_path["cached.csv"]["storage_origin"], "school_local_mirror")
        self.assertEqual(by_path["remote-only.csv"]["storage_origin"], "huggingface")

    def test_iccat_yearc_alias(self):
        path = self.root / "iccat.csv"
        path.write_text(
            "YearC,effort\n2023,2\n2023,3\n2024,4\n",
            encoding="utf-8",
        )
        result = fisheries_hf.analyze_and_export(
            path,
            dataset_path="ICCAT/processed/EFFDIS_LL2000-2024.csv",
            revision="b" * 40,
            repository="tuna",
        )
        self.assertEqual(result["time_range"], {"min_year": 2023, "max_year": 2024})
        self.assertEqual(result["annual_summary"][0]["year"], 2023)
        self.assertEqual(result["annual_summary"][0]["sum_effort"], 5.0)
        self.assertEqual(len(result["exports"]), 3)


if __name__ == "__main__":
    unittest.main()