From f4a3bafb43bcc3580d7a30cbfff13752f121833c Mon Sep 17 00:00:00 2001
From: Lee Kelvin <lkelvin@astro.princeton.edu>
Date: Tue, 18 Feb 2025 06:57:06 -0800
Subject: [PATCH] Add new unit tests for show_dataset_types=True

---
 tests/test_cliCmdQueryCollections.py | 61 ++++++++++++++++++++++++++++
 1 file changed, 61 insertions(+)

diff --git a/tests/test_cliCmdQueryCollections.py b/tests/test_cliCmdQueryCollections.py
index e9a7eb15d9..cce87da71d 100644
--- a/tests/test_cliCmdQueryCollections.py
+++ b/tests/test_cliCmdQueryCollections.py
@@ -93,6 +93,34 @@ def test_all(self):
                 exclude_dataset_types=("*_config", "*_log", "*_metadata", "packages"),
             ),
         )
+
+    def test_show_dataset_types(self):
+        """Test the --show-dataset-types option."""
+        self.run_test(
+            ["query-collections", "here", "--chains", "TABLE", "--show-dataset-types"],
+            self.makeExpected(
+                repo="here",
+                show_dataset_types=True,
+                exclude_dataset_types=("*_config", "*_log", "*_metadata", "packages"),
+            ),
+        )
+
+    def test_exclude_dataset_types(self):
+        """Test the --exclude-dataset-types option."""
+        self.run_test(
+            [
+                "query-collections",
+                "here",
+                "--chains",
+                "TABLE",
+                "--show-dataset-types",
+                "--exclude-dataset-types",
+                "flat",
+            ],
+            self.makeExpected(
+                repo="here",
+                show_dataset_types=True,
+                exclude_dataset_types=("flat",),
             ),
         )
 
@@ -230,6 +258,39 @@ def testChained(self):
             )
             self.assertAstropyTablesEqual(table, expected)
 
+            # Test table with show_dataset_types == True
+            table = queryCollections(
+                "here",
+                glob=(),
+                collection_type=CollectionType.all(),
+                chains="TREE",
+                show_dataset_types=True,
+            )
+            expected = Table(
+                array(
+                    (
+                        ("calibration1", "CALIBRATION", ""),
+                        ("chain1", "CHAINED", ""),
+                        ("  tag1", "TAGGED", ""),
+                        ("  run1", "RUN", ""),
+                        ("  chain2", "CHAINED", ""),
+                        ("    calibration1", "CALIBRATION", ""),
+                        ("    run1", "RUN", ""),
+                        ("chain2", "CHAINED", ""),
+                        ("  calibration1", "CALIBRATION", ""),
+                        ("  run1", "RUN", ""),
+                        ("imported_g", "RUN", "bias"),
+                        ("", "", "flat"),
+                        ("imported_r", "RUN", "bias"),
+                        ("", "", "flat"),
+                        ("run1", "RUN", ""),
+                        ("tag1", "TAGGED", ""),
+                    )
+                ),
+                names=("Name", "Type", "Dataset Types"),
+            )
+            self.assertAstropyTablesEqual(table, expected)
+
             result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "TABLE"])
             self.assertEqual(result.exit_code, 0, clickResultMsg(result))
             expected = Table(