From d242594e0f8550841fc8c0916f6ee64ebda75649 Mon Sep 17 00:00:00 2001
From: delucchi-cmu <delucchi@andrew.cmu.edu>
Date: Wed, 13 Dec 2023 09:46:54 -0500
Subject: [PATCH] Add another test for creating index table.

---
 tests/hipscat_import/index/test_run_index.py | 43 ++++++++++++++++++++
 1 file changed, 43 insertions(+)

diff --git a/tests/hipscat_import/index/test_run_index.py b/tests/hipscat_import/index/test_run_index.py
index f969fc1d..80190d10 100644
--- a/tests/hipscat_import/index/test_run_index.py
+++ b/tests/hipscat_import/index/test_run_index.py
@@ -65,3 +65,46 @@ def test_run_index(
 
     schema = pq.read_metadata(os.path.join(args.catalog_path, "_common_metadata")).schema.to_arrow_schema()
     assert schema.equals(basic_index_parquet_schema, check_metadata=False)
+
+
+@pytest.mark.dask
+def test_run_index_on_source(
+    small_sky_source_catalog,
+    tmp_path,
+):
+    """Test appropriate metadata is written, when primary catalog covers multiple pixels."""
+
+    args = IndexArguments(
+        input_catalog_path=small_sky_source_catalog,
+        indexing_column="source_id",
+        output_path=tmp_path,
+        output_artifact_name="small_sky_source_id_index",
+        overwrite=True,
+        progress_bar=False,
+    )
+    runner.run(args)
+
+    # Check that the catalog metadata file exists
+    catalog = Dataset.read_from_hipscat(args.catalog_path)
+    assert catalog.on_disk
+    assert catalog.catalog_path == args.catalog_path
+
+    basic_index_parquet_schema = pa.schema(
+        [
+            pa.field("_hipscat_index", pa.uint64()),
+            pa.field("Norder", pa.uint8()),
+            pa.field("Dir", pa.uint64()),
+            pa.field("Npix", pa.uint64()),
+            pa.field("source_id", pa.int64()),
+        ]
+    )
+
+    outfile = os.path.join(args.catalog_path, "index", "part.0.parquet")
+    schema = pq.read_metadata(outfile).schema.to_arrow_schema()
+    assert schema.equals(basic_index_parquet_schema, check_metadata=False)
+
+    schema = pq.read_metadata(os.path.join(args.catalog_path, "_metadata")).schema.to_arrow_schema()
+    assert schema.equals(basic_index_parquet_schema, check_metadata=False)
+
+    schema = pq.read_metadata(os.path.join(args.catalog_path, "_common_metadata")).schema.to_arrow_schema()
+    assert schema.equals(basic_index_parquet_schema, check_metadata=False)