diff --git a/backend/core/test/test_product_registry.py b/backend/core/test/test_product_registry.py index 8f12888..77b7f02 100644 --- a/backend/core/test/test_product_registry.py +++ b/backend/core/test/test_product_registry.py @@ -9,7 +9,7 @@ from unittest import mock import pandas as pd -from core.models import Product, ProductFile, ProductType, Release +from core.models import Product, ProductContent, ProductFile, ProductType, Release from core.product_handle import ProductHandle from core.product_steps import RegistryProduct from core.table_data_collector import MainTableDataCollector @@ -134,7 +134,7 @@ def create_multi_file_archive(self, extension="zip"): archive.add(file_b, arcname=file_b.name) return archive_path - def create_fake_hats_archive(self, properties_filename="hats.properties"): + def create_fake_hats_archive(self, properties_filename="hats.properties", data=None): temp_dir = Path(tempfile.mkdtemp(prefix="pz_hats_archive_")) hats_root = temp_dir / "mock_hats" dataset_dir = hats_root / "dataset" / "Norder=0" / "Npix=0" @@ -144,7 +144,14 @@ def create_fake_hats_archive(self, properties_filename="hats.properties"): "catalog_name=mock_hats\ndataproduct_type=object\n", encoding="utf-8", ) - (dataset_dir / "part-0.parquet").write_text("", encoding="utf-8") + if data is None: + data = pd.DataFrame( + [ + {"ra": 0.0, "dec": 0.0, "z": 0.0}, + {"ra": 1.0, "dec": -1.0, "z": 0.1}, + ] + ) + data.to_parquet(dataset_dir / "part-0.parquet", index=False) archive_path = temp_dir / "mock_hats.zip" with zipfile.ZipFile(archive_path, "w") as archive: @@ -176,8 +183,15 @@ def create_fake_hats_collection_archive(self): "catalog_name=mock_margin\ndataproduct_type=margin\n", encoding="utf-8", ) - (catalog_dataset_dir / "part-0.parquet").write_text("", encoding="utf-8") - (margin_dataset_dir / "part-0.parquet").write_text("", encoding="utf-8") + pd.DataFrame( + [ + {"ra": 0.0, "dec": 0.0, "z": 0.0}, + {"ra": 1.0, "dec": -1.0, "z": 0.1}, + ] + ).to_parquet(catalog_dataset_dir / "part-0.parquet", index=False) + pd.DataFrame([{"margin_id": 1}]).to_parquet( + margin_dataset_dir / "part-0.parquet", index=False + ) archive_path = temp_dir / "mock_collection.zip" with zipfile.ZipFile(archive_path, "w") as archive: @@ -187,19 +201,43 @@ def create_fake_hats_collection_archive(self): archive.write(file_path, arcname=arcname) return archive_path - def create_fake_lsdb_module(self, n_rows=12, columns_as_index=False): + def create_fake_lsdb_module( + self, + n_rows=12, + columns_as_index=False, + columns=None, + all_columns=None, + preview_columns=None, + ): + catalog_columns = columns or ["ra", "dec", "z"] + catalog_all_columns = ( + all_columns if all_columns is not None else catalog_columns + ) + catalog_preview_columns = preview_columns or catalog_columns + class FakeCatalog: - columns = pd.Index(["ra", "dec", "z"]) if columns_as_index else ["ra", "dec", "z"] - all_columns = ["ra", "dec", "z"] + columns = ( + pd.Index(catalog_columns) if columns_as_index else catalog_columns + ) + all_columns = catalog_all_columns def __len__(self): return n_rows def head(self, n=5): - rows = [ - {"ra": float(i), "dec": float(i) * -1.0, "z": float(i) * 0.1} - for i in range(n_rows) - ] + rows = [] + for i in range(n_rows): + row = {} + for column in catalog_preview_columns: + if column == "ra": + row[column] = float(i) + elif column == "dec": + row[column] = float(i) * -1.0 + elif column == "z": + row[column] = float(i) * 0.1 + else: + row[column] = f"{column}_{i}" + rows.append(row) return pd.DataFrame(rows).head(n) fake_module = types.ModuleType("lsdb") @@ -428,6 +466,36 @@ def test_registry_handles_lsdb_columns_as_pandas_index(self): self.assertEqual(response.status_code, 200) + def test_registry_uses_lsdb_columns_for_hats_content_columns(self): + product = self.create_product(product_type_name="validation_results") + expected_columns = ["object_id", "ra", "dec", "z", "quality_flag"] + archive_path = self.create_fake_hats_archive() + self.upload_main_file_from_path(product, archive_path) + url = reverse("products-registry", kwargs={"pk": product.pk}) + fake_lsdb = self.create_fake_lsdb_module( + n_rows=1, + columns=expected_columns, + all_columns=expected_columns, + preview_columns=expected_columns, + ) + + with mock.patch.dict(sys.modules, {"lsdb": fake_lsdb}): + response = self.client.post(url) + + self.assertEqual(response.status_code, 200) + registered_columns = list( + ProductContent.objects.filter(product=product) + .order_by("order") + .values_list("column_name", flat=True) + ) + self.assertEqual(registered_columns, expected_columns) + + preview_path = Path( + settings.MEDIA_ROOT, product.path, RegistryProduct.TABLE_PREVIEW_FILENAME + ) + payload = json.loads(preview_path.read_text(encoding="utf-8")) + self.assertEqual(payload["columns"], expected_columns) + def test_registry_without_columns(self): # Cria um novo produto.