diff --git a/src/worldcereal/rdm_api/rdm_interaction.py b/src/worldcereal/rdm_api/rdm_interaction.py index a86348e..1eae901 100644 --- a/src/worldcereal/rdm_api/rdm_interaction.py +++ b/src/worldcereal/rdm_api/rdm_interaction.py @@ -398,6 +398,7 @@ def download_samples( ------- gpd.GeoDataFrame A GeoDataFrame containing the extracted samples. + For each sample, the collection ID is automatically included. """ # Determine which collections need to be queried if they are not specified diff --git a/tests/worldcerealtests/test_rdm_interaction.py b/tests/worldcerealtests/test_rdm_interaction.py index d85e6a3..995ac7c 100644 --- a/tests/worldcerealtests/test_rdm_interaction.py +++ b/tests/worldcerealtests/test_rdm_interaction.py @@ -100,7 +100,7 @@ def test_download_samples( } ), ] - mock_get_download_urls.return_value = [file_path] + mock_get_download_urls.return_value = [str(file_path)] interaction = RdmInteraction() result_gdf = interaction.download_samples( @@ -112,7 +112,13 @@ def test_download_samples( ) # Check that col3 and valid_time indeed not included - assert result_gdf.columns.tolist() == ["col1", "col2", "geometry"] + # collection_id is automatically added + assert result_gdf.columns.tolist() == [ + "col1", + "col2", + "collection_id", + "geometry", + ] # Check that the third up till last geometry are not included # third and fourth are outside the spatiotemporal extent