From 4c2564d4bf9fe6f8e14ee4e1b7a5fad590120d2d Mon Sep 17 00:00:00 2001
From: Dougal Dobie <ddobie@users.noreply.github.com>
Date: Sat, 23 Sep 2023 14:02:49 +1000
Subject: [PATCH] Exit nicely when invalid survey is requested for SkyView
 contour plot (#497)

* Initial fix

* Add exception test

* Updated changelog

* PEP8

* PEP8
---
 CHANGELOG.md         |  2 ++
 tests/test_source.py | 36 +++++++++++++++++++++++++++++++++++-
 vasttools/source.py  |  7 +++++++
 3 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 83548e4f..827dc7f7 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,12 +10,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 #### Changed
 
+- Exit nicely when invalid survey is requested for SkyView contour plot [#497](https://github.com/askap-vast/vast-tools/pull/497)
 #### Fixed
 
 #### Removed
 
 #### List of PRs
 
+- [#497](https://github.com/askap-vast/vast-tools/pull/497): feat: Exit nicely when invalid survey is requested for SkyView contour plot
 
 ## [3.0.0](https://github.com/askap-vast/vast-tools/releases/v3.0.0) (2023-09-02)
 
diff --git a/tests/test_source.py b/tests/test_source.py
index bedabb99..cd17b5b8 100644
--- a/tests/test_source.py
+++ b/tests/test_source.py
@@ -892,11 +892,45 @@ def test_skyview_contour_plot(
             return_value=[dummy_fits]
         )
 
-        result = source.skyview_contour_plot(0, 'suveycode')
+        result = source.skyview_contour_plot(0, 'DSS2 Blue')
 
         assert isinstance(result, Figure)
         plt.close(result)
 
+    @pytest.mark.parametrize("pipeline", [False, True])
+    def test_skyview_contour_plot_survey_fail(
+        self,
+        pipeline: bool,
+        source_instance: vts.Source,
+        dummy_fits: fits.HDUList,
+        mocker: MockerFixture
+    ) -> None:
+        """
+        Tests the skyview_contour_plot method.
+
+        Parametrized for pipeline and query source.
+
+        Args:
+            pipeline: If 'True' then the Source is initialised as a
+                pipeline source.
+            source_instance: The pytest source_instance fixture.
+            dummy_fits: The pytest fixture dummy fits.
+            mocker: The pytest-mock mocker object.
+
+        Returns:
+            None
+        """
+        source = source_instance(pipeline=pipeline, add_cutout_data=True)
+
+        mocker_skyview = mocker.patch(
+            'vasttools.source.SkyView.get_images',
+            return_value=[dummy_fits]
+        )
+        with pytest.raises(ValueError) as excinfo:
+            source.skyview_contour_plot(0, 'this-is-not-a-survey')
+
+        assert str(excinfo.value).endswith('not a valid SkyView survey name')
+
     @pytest.mark.parametrize("pipeline", [False, True])
     def test_write_ann(
         self,
diff --git a/vasttools/source.py b/vasttools/source.py
index b15d2030..d91ebf4c 100644
--- a/vasttools/source.py
+++ b/vasttools/source.py
@@ -1651,6 +1651,7 @@ def skyview_contour_plot(
 
         Raises:
             ValueError: If the index is out of range.
+            ValueError: If the requested survey is not valid.
         """
 
         if (self._cutouts_got is False) or (force):
@@ -1658,6 +1659,12 @@ def skyview_contour_plot(
 
         size = self._size
 
+        surveys = list(SkyView.survey_dict.values())
+        survey_list = [item for sublist in surveys for item in sublist]
+
+        if survey not in survey_list:
+            raise ValueError(f"{survey} is not a valid SkyView survey name")
+
         if index > len(self.measurements):
             raise ValueError(f"Cannot access {index}th measurement.")
             return