From ea36f51fdc8d519c7feea95cf5755c4ebce3e561 Mon Sep 17 00:00:00 2001
From: Duy Nguyen <duytnguyendtn.open@gmail.com>
Date: Tue, 23 Jul 2024 11:59:59 -0400
Subject: [PATCH] Add custom table header test

---
 .../tests/test_vo_imviz.py                    | 48 +++++++++++++++----
 1 file changed, 38 insertions(+), 10 deletions(-)

diff --git a/jdaviz/configs/default/plugins/virtual_observatory/tests/test_vo_imviz.py b/jdaviz/configs/default/plugins/virtual_observatory/tests/test_vo_imviz.py
index 2f62758f5a..a117029520 100644
--- a/jdaviz/configs/default/plugins/virtual_observatory/tests/test_vo_imviz.py
+++ b/jdaviz/configs/default/plugins/virtual_observatory/tests/test_vo_imviz.py
@@ -4,6 +4,15 @@
 from jdaviz.configs.default.plugins.virtual_observatory.vo_plugin import vo_plugin_label
 
 
+class fake_siaresult:
+    def __init__(self, attrs):
+        for attr, value in attrs.items():
+            self.__setattr__(attr, value)
+
+    def getdataurl(self):
+        return "Fake URL"
+
+
 # TODO: Update all _obj calls to formal API calls once Plugin API is available
 class TestVOImvizLocal(BaseImviz_WCS_WCS):
     def test_autocenter_coords(self):
@@ -25,17 +34,14 @@ def test_autocenter_coords(self):
         assert vo_plugin.source == "337.51924057481 -20.83208305686149"
 
     def test_populate_table_default_headers(self):
-        class fake_siaresult:
-            def __init__(self):
-                self.title = "Fake Title"
-                self.instr = "Fake Instrument"
-                self.dateobs = "Fake Dateobs"
-
-            def getdataurl(self):
-                return "Fake URL"
-
         vo_plugin = self.imviz.plugins[vo_plugin_label]._obj
-        fake_result = fake_siaresult()
+        fake_result = fake_siaresult(
+            {
+                "title": "Fake Title",
+                "instr": "Fake Instrument",
+                "dateobs": "Fake Dateobs",
+            }
+        )
         vo_plugin._populate_table([fake_result])
 
         assert vo_plugin.table.items == [
@@ -47,6 +53,28 @@ def getdataurl(self):
             }
         ]
 
+    def test_populate_table_customm_headers(self):
+        vo_plugin = self.imviz.plugins[vo_plugin_label]._obj
+        fake_result = fake_siaresult(
+            {
+                "attrA": "Field A",
+                "attrB": "Field B",
+                "attrC": "Field C",
+            }
+        )
+        vo_plugin._populate_table(
+            [fake_result], {"Title A": "attrA", "Title B": "attrB", "Title C": "attrC"}
+        )
+
+        assert vo_plugin.table.items == [
+            {
+                "URL": fake_result.getdataurl(),
+                "Title A": fake_result.attrA,
+                "Title B": fake_result.attrB,
+                "Title C": fake_result.attrC,
+            }
+        ]
+
 
 @pytest.mark.remote_data
 class TestVOImvizRemote: