earthcatalog 0.2.0__py3-none-any.whl

earthcatalog/tests/test_workers.py

@@ -0,0 +1,653 @@
+# test_workers.py
+"""Tests for workers module - serializable worker functions.
+
+This module tests:
+- process_url_batch: URL batch processing
+- consolidate_partition: Partition consolidation
+- Result types: DownloadResult, ConsolidationResult
+"""
+
+import tempfile
+from pathlib import Path
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from earthcatalog.ingestion_pipeline import ProcessingConfig
+from earthcatalog.statistics import IngestionStatistics
+from earthcatalog.workers import (
+    ConsolidationResult,
+    DownloadResult,
+    _download_stac_item,
+    _get_partition_key,
+    consolidate_partition,
+    process_url_batch,
+)
+
+
+class TestDownloadResult:
+    """Test DownloadResult dataclass."""
+
+    def test_to_dict(self):
+        """to_dict should serialize result."""
+        stats = IngestionStatistics()
+        stats.record_url_processed(success=True)
+        result = DownloadResult(
+            shards=[{"shard_path": "/path/to/shard.parquet"}],
+            stats=stats,
+            failed_urls=["http://failed.com"],
+        )
+
+        data = result.to_dict()
+
+        assert data["shards"] == [{"shard_path": "/path/to/shard.parquet"}]
+        assert data["failed_urls"] == ["http://failed.com"]
+        # Stats contains url counts
+        assert "urls_processed" in data["stats"]
+
+    def test_from_dict(self):
+        """from_dict should deserialize result."""
+        data = {
+            "shards": [{"shard_path": "/path/to/shard.parquet"}],
+            "stats": {"items_processed": 10, "items_failed": 2},
+            "failed_urls": ["http://failed.com"],
+        }
+
+        result = DownloadResult.from_dict(data)
+
+        assert len(result.shards) == 1
+        assert len(result.failed_urls) == 1
+
+
+class TestConsolidationResult:
+    """Test ConsolidationResult dataclass."""
+
+    def test_to_dict(self):
+        """to_dict should serialize result."""
+        result = ConsolidationResult(
+            partition_key="h3_82/2024/01",
+            item_count=100,
+            existing_count=50,
+            new_count=50,
+            duplicates_removed=5,
+            final_path="/catalog/h3_82/2024/01/data.parquet",
+        )
+
+        data = result.to_dict()
+
+        assert data["partition_key"] == "h3_82/2024/01"
+        assert data["item_count"] == 100
+        assert data["existing_count"] == 50
+        assert data["new_count"] == 50
+        assert data["duplicates_removed"] == 5
+
+
+class TestDownloadStacItem:
+    """Test _download_stac_item helper function."""
+
+    def test_successful_download(self):
+        """Should return parsed JSON on successful download."""
+        mock_hook = MagicMock()
+        mock_hook.fetch.return_value = {"id": "test-item", "type": "Feature"}
+
+        with patch("earthcatalog.workers._get_stac_hook", return_value=mock_hook):
+            result = _download_stac_item("http://example.com/item.json")
+
+        assert result is not None
+        assert result["id"] == "test-item"
+        mock_hook.fetch.assert_called_once()
+
+    def test_failed_download_returns_none(self):
+        """Should return None after all retries fail."""
+        mock_hook = MagicMock()
+        mock_hook.fetch.return_value = None
+
+        with patch("earthcatalog.workers._get_stac_hook", return_value=mock_hook):
+            result = _download_stac_item("http://example.com/item.json", retry_attempts=1)
+
+        assert result is None
+
+
+class TestGetPartitionKey:
+    """Test _get_partition_key helper function."""
+
+    def test_partition_key_with_datetime(self):
+        """Should generate correct partition key with datetime."""
+        from earthcatalog.grid_systems import get_grid_system
+
+        item = {
+            "id": "test-item",
+            "geometry": {"type": "Point", "coordinates": [-122.0, 37.0]},
+            "properties": {"datetime": "2024-06-15T12:00:00Z"},
+        }
+        grid = get_grid_system("h3", resolution=2)
+
+        key = _get_partition_key(
+            item,
+            grid_resolver=grid,
+            temporal_bin="month",
+            enable_global=False,
+            global_threshold=50,
+        )
+
+        assert "2024" in key
+        assert "06" in key
+
+    def test_partition_key_with_year_bin(self):
+        """Should generate year-level partition key."""
+        from earthcatalog.grid_systems import get_grid_system
+
+        item = {
+            "id": "test-item",
+            "geometry": {"type": "Point", "coordinates": [-122.0, 37.0]},
+            "properties": {"datetime": "2024-06-15T12:00:00Z"},
+        }
+        grid = get_grid_system("h3", resolution=2)
+
+        key = _get_partition_key(
+            item,
+            grid_resolver=grid,
+            temporal_bin="year",
+            enable_global=False,
+            global_threshold=50,
+        )
+
+        assert "2024" in key
+        # Month should not be in year-level binning
+        assert key.count("/") == 1  # grid_cell/year
+
+
+class TestProcessUrlBatch:
+    """Test process_url_batch function."""
+
+    @pytest.fixture
+    def temp_dir(self):
+        """Create a temporary directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield tmpdir
+
+    @pytest.fixture
+    def config_dict(self, temp_dir):
+        """Create a config dict for testing."""
+        config = ProcessingConfig(
+            input_file="test.parquet",
+            output_catalog=f"{temp_dir}/catalog",
+            scratch_location=f"{temp_dir}/scratch",
+            enable_concurrent_http=False,  # Disable async for testing
+            batch_size=10,
+        )
+        return config.to_dict()
+
+    def test_process_empty_batch(self, config_dict):
+        """Should handle empty URL list."""
+        result = process_url_batch(
+            urls=[],
+            worker_id=0,
+            config_dict=config_dict,
+            job_id="test-job",
+        )
+
+        assert result["shards"] == []
+        assert result["failed_urls"] == []
+
+    def test_process_batch_with_failures(self, config_dict):
+        """Should record failed URLs."""
+        mock_hook = MagicMock()
+        mock_hook.fetch.return_value = None  # Simulate failed fetch
+        mock_hook.fetch_batch.return_value = [None, None]  # Batch also fails
+
+        with patch("earthcatalog.workers._get_stac_hook", return_value=mock_hook):
+            result = process_url_batch(
+                urls=["http://bad1.com", "http://bad2.com"],
+                worker_id=0,
+                config_dict=config_dict,
+                job_id="test-job",
+            )
+
+        assert len(result["failed_urls"]) == 2
+
+    @pytest.mark.skip(reason="Requires complex mocking of STAC engine serialization")
+    def test_process_batch_successful(self, config_dict):
+        """Should process URLs and write shards."""
+        mock_item = {
+            "id": "test-item",
+            "type": "Feature",
+            "geometry": {"type": "Point", "coordinates": [-122.0, 37.0]},
+            "bbox": [-122.0, 37.0, -122.0, 37.0],
+            "properties": {"datetime": "2024-06-15T12:00:00Z"},
+            "stac_version": "1.0.0",
+            "links": [],
+        }
+
+        mock_response = MagicMock()
+        mock_response.json.return_value = mock_item
+        mock_response.raise_for_status = MagicMock()
+
+        with patch("earthcatalog.workers.requests.get", return_value=mock_response):
+            result = process_url_batch(
+                urls=["http://example.com/item1.json"],
+                worker_id=0,
+                config_dict=config_dict,
+                job_id="test-job",
+            )
+
+        assert len(result["shards"]) >= 1
+        assert result["failed_urls"] == []
+
+
+class TestConsolidatePartition:
+    """Test consolidate_partition function."""
+
+    @pytest.fixture
+    def temp_dir(self):
+        """Create a temporary directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield tmpdir
+
+    @pytest.fixture
+    def config_dict(self, temp_dir):
+        """Create a config dict for testing."""
+        config = ProcessingConfig(
+            input_file="test.parquet",
+            output_catalog=f"{temp_dir}/catalog",
+            scratch_location=f"{temp_dir}/scratch",
+        )
+        return config.to_dict()
+
+    def test_consolidate_empty_shards(self, config_dict):
+        """Should handle empty shard list."""
+        result = consolidate_partition(
+            partition_key="h3_82/2024/01",
+            shard_paths=[],
+            config_dict=config_dict,
+        )
+
+        assert result["partition_key"] == "h3_82/2024/01"
+        assert result["item_count"] == 0
+
+    def test_consolidate_with_shards(self, config_dict, temp_dir):
+        """Should consolidate shards into final partition."""
+        import geopandas as gpd
+        from shapely.geometry import Point
+
+        # Create a test shard
+        scratch_dir = Path(temp_dir) / "scratch" / "shards"
+        scratch_dir.mkdir(parents=True)
+
+        shard_path = str(scratch_dir / "test_shard.parquet")
+        gdf = gpd.GeoDataFrame(
+            {
+                "id": ["item1", "item2"],
+                "datetime": ["2024-06-15T12:00:00Z", "2024-06-16T12:00:00Z"],
+            },
+            geometry=gpd.GeoSeries([Point(-122, 37), Point(-122, 37)]),
+        )
+        gdf.to_parquet(shard_path)
+
+        result = consolidate_partition(
+            partition_key="h3_82/2024/06",
+            shard_paths=[shard_path],
+            config_dict=config_dict,
+        )
+
+        assert result["partition_key"] == "h3_82/2024/06"
+        assert result["item_count"] == 2
+        assert Path(result["final_path"]).exists()
+
+    def test_consolidate_deduplicates_by_datetime(self, config_dict, temp_dir):
+        """Should keep newer item when deduplicating by ID."""
+        import geopandas as gpd
+        from shapely.geometry import Point
+
+        scratch_dir = Path(temp_dir) / "scratch" / "shards"
+        scratch_dir.mkdir(parents=True)
+
+        # Create two shards with same ID but different datetimes
+        shard1_path = str(scratch_dir / "shard1.parquet")
+        gdf1 = gpd.GeoDataFrame(
+            {
+                "id": ["item1"],
+                "datetime": ["2024-06-15T12:00:00Z"],  # Older
+            },
+            geometry=gpd.GeoSeries([Point(-122, 37)]),
+        )
+        gdf1.to_parquet(shard1_path)
+
+        shard2_path = str(scratch_dir / "shard2.parquet")
+        gdf2 = gpd.GeoDataFrame(
+            {
+                "id": ["item1"],
+                "datetime": ["2024-06-16T12:00:00Z"],  # Newer
+            },
+            geometry=gpd.GeoSeries([Point(-122, 37)]),
+        )
+        gdf2.to_parquet(shard2_path)
+
+        result = consolidate_partition(
+            partition_key="h3_82/2024/06",
+            shard_paths=[shard1_path, shard2_path],
+            config_dict=config_dict,
+        )
+
+        assert result["item_count"] == 1  # Deduplicated to 1
+        assert result["duplicates_removed"] == 1
+
+        # Verify the newer item was kept
+        final_gdf = gpd.read_parquet(result["final_path"])
+        assert len(final_gdf) == 1
+        assert "2024-06-16" in final_gdf.iloc[0]["datetime"]
+
+
+class TestPickleability:
+    """Test that worker functions can be pickled for Dask."""
+
+    def test_process_url_batch_is_pickleable(self):
+        """process_url_batch should be pickleable."""
+        import pickle
+
+        # Should not raise
+        pickled = pickle.dumps(process_url_batch)
+        restored = pickle.loads(pickled)
+        assert callable(restored)
+
+    def test_consolidate_partition_is_pickleable(self):
+        """consolidate_partition should be pickleable."""
+        import pickle
+
+        pickled = pickle.dumps(consolidate_partition)
+        restored = pickle.loads(pickled)
+        assert callable(restored)
+
+
+class TestDaskIntegration:
+    """Tests for Dask distributed integration with serializable workers.
+
+    These tests verify that the worker functions can be properly used
+    with Dask's distributed computing model.
+    """
+
+    @pytest.fixture
+    def temp_dir(self):
+        """Create a temporary directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield tmpdir
+
+    @pytest.fixture
+    def config_dict(self, temp_dir):
+        """Create a config dict for testing."""
+        config = ProcessingConfig(
+            input_file="test.parquet",
+            output_catalog=f"{temp_dir}/catalog",
+            scratch_location=f"{temp_dir}/scratch",
+            enable_concurrent_http=False,
+            batch_size=10,
+        )
+        return config.to_dict()
+
+    def test_config_dict_roundtrip(self, config_dict):
+        """Config dict should serialize and deserialize correctly."""
+        restored = ProcessingConfig.from_dict(config_dict)
+
+        assert restored.output_catalog == config_dict["output_catalog"]
+        assert restored.scratch_location == config_dict["scratch_location"]
+        assert restored.enable_concurrent_http == config_dict["enable_concurrent_http"]
+
+    def test_process_url_batch_accepts_config_dict(self, config_dict):
+        """process_url_batch should accept config_dict parameter."""
+        # Should not raise
+        result = process_url_batch(
+            urls=[],
+            worker_id=0,
+            config_dict=config_dict,
+            job_id="test-job-123",
+        )
+
+        assert "shards" in result
+        assert "stats" in result
+        assert "failed_urls" in result
+
+    def test_consolidate_partition_accepts_config_dict(self, config_dict):
+        """consolidate_partition should accept config_dict parameter."""
+        result = consolidate_partition(
+            partition_key="test/2024/01",
+            shard_paths=[],
+            config_dict=config_dict,
+        )
+
+        assert "partition_key" in result
+        assert result["partition_key"] == "test/2024/01"
+        assert result["item_count"] == 0
+
+    def test_worker_functions_are_module_level(self):
+        """Worker functions should be importable at module level for Dask."""
+        from earthcatalog import workers
+
+        # These should be module-level functions, not methods
+        assert hasattr(workers, "process_url_batch")
+        assert hasattr(workers, "consolidate_partition")
+        assert callable(workers.process_url_batch)
+        assert callable(workers.consolidate_partition)
+
+    def test_worker_functions_have_correct_signature(self):
+        """Worker functions should have correct signature for Dask usage."""
+        import inspect
+
+        # process_url_batch signature
+        sig = inspect.signature(process_url_batch)
+        params = list(sig.parameters.keys())
+        assert "urls" in params
+        assert "worker_id" in params
+        assert "config_dict" in params
+        assert "job_id" in params
+
+        # consolidate_partition signature
+        sig = inspect.signature(consolidate_partition)
+        params = list(sig.parameters.keys())
+        assert "partition_key" in params
+        assert "shard_paths" in params
+        assert "config_dict" in params
+
+    def test_result_serialization_for_dask(self, config_dict):
+        """Results should be serializable for Dask transmission."""
+        import pickle
+
+        result = process_url_batch(
+            urls=[],
+            worker_id=0,
+            config_dict=config_dict,
+            job_id="test",
+        )
+
+        # Should be pickleable (required for Dask)
+        pickled = pickle.dumps(result)
+        restored = pickle.loads(pickled)
+
+        assert restored["shards"] == result["shards"]
+        assert restored["failed_urls"] == result["failed_urls"]
+
+    def test_consolidation_result_serialization(self, config_dict):
+        """Consolidation results should be serializable."""
+        import pickle
+
+        result = consolidate_partition(
+            partition_key="h3_82/2024/01",
+            shard_paths=[],
+            config_dict=config_dict,
+        )
+
+        pickled = pickle.dumps(result)
+        restored = pickle.loads(pickled)
+
+        assert restored["partition_key"] == result["partition_key"]
+        assert restored["item_count"] == result["item_count"]
+
+
+class TestDaskDistributedProcessorIntegration:
+    """Tests for DaskDistributedProcessor using workers.py functions.
+
+    These tests mock Dask to verify the correct functions are called.
+    """
+
+    @pytest.fixture
+    def mock_dask_client(self):
+        """Create a mock Dask client."""
+        mock_client = MagicMock()
+        mock_client.dashboard_link = "http://localhost:8787"
+        mock_client.submit.return_value = MagicMock()
+        mock_client.gather.return_value = [{"shards": [], "stats": {}, "failed_urls": []}]
+        return mock_client
+
+    def test_dask_processor_uses_workers_process_url_batch(self, mock_dask_client):
+        """DaskDistributedProcessor should use workers.process_url_batch when config_dict provided."""
+        # Create mock dd module
+        mock_dd = MagicMock()
+        mock_dd.Client.return_value = mock_dask_client
+
+        with patch.dict("sys.modules", {"dask.distributed": mock_dd}):
+            # Need to reimport to pick up mocked module
+            from importlib import reload
+
+            from earthcatalog import ingestion_pipeline
+
+            reload(ingestion_pipeline)
+
+            processor = ingestion_pipeline.DaskDistributedProcessor()
+
+            config_dict = {
+                "input_file": "test.parquet",
+                "output_catalog": "./out",
+                "scratch_location": "./scratch",
+            }
+
+            # Call process_urls with config_dict
+            processor.process_urls(
+                url_chunks=[["url1", "url2"]],
+                process_fn=lambda x, y: None,  # Should be ignored
+                config_dict=config_dict,
+                job_id="test-job",
+            )
+
+            # Verify the submit call used process_url_batch
+            assert mock_dask_client.submit.called
+            call_args = mock_dask_client.submit.call_args
+            # First positional arg should be the function
+            called_func = call_args[0][0] if call_args[0] else call_args[1].get("func")
+            # The function name should be process_url_batch
+            assert called_func.__name__ == "process_url_batch"
+
+    def test_dask_processor_uses_workers_consolidate_partition(self, mock_dask_client):
+        """DaskDistributedProcessor should use workers.consolidate_partition when config_dict provided."""
+        mock_dask_client.gather.return_value = [{"partition_key": "test", "item_count": 0}]
+        mock_dd = MagicMock()
+        mock_dd.Client.return_value = mock_dask_client
+
+        with patch.dict("sys.modules", {"dask.distributed": mock_dd}):
+            from importlib import reload
+
+            from earthcatalog import ingestion_pipeline
+
+            reload(ingestion_pipeline)
+
+            processor = ingestion_pipeline.DaskDistributedProcessor()
+
+            config_dict = {
+                "input_file": "test.parquet",
+                "output_catalog": "./out",
+                "scratch_location": "./scratch",
+            }
+
+            processor.consolidate_shards(
+                partition_items=[("test/2024/01", ["/path/shard1.parquet"])],
+                consolidate_fn=lambda x, y: None,  # Should be ignored
+                config_dict=config_dict,
+            )
+
+            assert mock_dask_client.submit.called
+            call_args = mock_dask_client.submit.call_args
+            called_func = call_args[0][0] if call_args[0] else call_args[1].get("func")
+            assert called_func.__name__ == "consolidate_partition"
+
+    def test_dask_processor_fallback_without_config_dict(self, mock_dask_client):
+        """DaskDistributedProcessor should fall back to old behavior without config_dict."""
+        mock_dask_client.gather.return_value = ["result1"]
+        mock_dd = MagicMock()
+        mock_dd.Client.return_value = mock_dask_client
+
+        with patch.dict("sys.modules", {"dask.distributed": mock_dd}):
+            from importlib import reload
+
+            from earthcatalog import ingestion_pipeline
+
+            reload(ingestion_pipeline)
+
+            processor = ingestion_pipeline.DaskDistributedProcessor()
+
+            custom_fn = MagicMock(return_value="custom_result")
+
+            # Call without config_dict
+            processor.process_urls(
+                url_chunks=[["url1"]],
+                process_fn=custom_fn,
+                # No config_dict - should use custom_fn
+            )
+
+            assert mock_dask_client.submit.called
+            call_args = mock_dask_client.submit.call_args
+            called_func = call_args[0][0]
+            # Should use the provided custom function, not process_url_batch
+            assert called_func == custom_fn
+
+
+class TestStacHookIntegrationWithWorkers:
+    """Test STAC hook integration with workers.py."""
+
+    @pytest.fixture
+    def temp_dir(self):
+        """Create a temporary directory."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            yield tmpdir
+
+    def test_hook_config_in_config_dict(self, temp_dir):
+        """stac_hook should be included in config_dict."""
+        config = ProcessingConfig(
+            input_file="test.parquet",
+            output_catalog=f"{temp_dir}/catalog",
+            scratch_location=f"{temp_dir}/scratch",
+            stac_hook="module:my_module:my_func",
+        )
+
+        config_dict = config.to_dict()
+        assert "stac_hook" in config_dict
+        assert config_dict["stac_hook"] == "module:my_module:my_func"
+
+    def test_hook_config_restored_from_dict(self, temp_dir):
+        """stac_hook should be restored from config_dict."""
+        config_dict = {
+            "input_file": "test.parquet",
+            "output_catalog": f"{temp_dir}/catalog",
+            "scratch_location": f"{temp_dir}/scratch",
+            "stac_hook": "script:/path/to/script.py",
+        }
+
+        config = ProcessingConfig.from_dict(config_dict)
+        assert config.stac_hook == "script:/path/to/script.py"
+
+    def test_default_hook_in_workers(self, temp_dir):
+        """Workers should use default hook when not specified."""
+        config = ProcessingConfig(
+            input_file="test.parquet",
+            output_catalog=f"{temp_dir}/catalog",
+            scratch_location=f"{temp_dir}/scratch",
+        )
+
+        config_dict = config.to_dict()
+
+        # process_url_batch should work with default hook
+        result = process_url_batch(
+            urls=[],
+            worker_id=0,
+            config_dict=config_dict,
+            job_id="test",
+        )
+
+        assert result["failed_urls"] == []