From f04a6f8f7caebd59bc2925de777511d687383ece Mon Sep 17 00:00:00 2001
From: Matt Craig <mattwcraig@gmail.com>
Date: Tue, 2 Jan 2024 12:17:43 -0600
Subject: [PATCH] Finish validators in exoplanet model

---
 stellarphot/settings/models.py            | 22 +++++++++++-------
 stellarphot/settings/tests/test_models.py | 28 ++++++++---------------
 2 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/stellarphot/settings/models.py b/stellarphot/settings/models.py
index d1f59b31..05938072 100644
--- a/stellarphot/settings/models.py
+++ b/stellarphot/settings/models.py
@@ -143,6 +143,9 @@ class Exoplanet(BaseModel):
     To create an `Exoplanet` object, you can pass in the epoch,
      period, identifier, coordinate, depth, and duration as keyword arguments:
 
+    >>> from astropy.time import Time
+    >>> from astropy.coordinates import SkyCoord
+    >>> from astropy import units as u
     >>> planet  = Exoplanet(epoch=Time(2455909.29280, format="jd"),
     ...                     period=1.21749 * u.day,
     ...                     identifier="KELT-1b",
@@ -170,25 +173,28 @@ class Config:
             QuantityType: lambda v: f"{v.value} {v.unit}",
             Time: lambda v: f"{v.value}",
         }
-    @validator
+
+    @validator("period")
     @classmethod
-    def validate_period(cls, values):
+    def validate_period(cls, value):
         """
         Checks that the period has physical units of time and raises an error if that is not true. 
         """
-        if u.get_physical_type(values["period"]) != "time":
+        if u.get_physical_type(value) != "time":
             raise ValueError(
-                f"Period does not have time units,"
-                f"currently has {values['period'].unit} units."
+                f"Period does not have time units," f"currently has {value.unit} units."
             )
+        return value
 
+    @validator("duration")
     @classmethod
-    def validate_duration(cls, values):
+    def validate_duration(cls, value):
         """
         Checks that the duration has physical units of time and raises an error if that is not true. 
         """
-        if u.get_physical_type(values["duration"]) != "time":
+        if u.get_physical_type(value) != "time":
             raise ValueError(
                 f"Duration does not have time units,"
-                f"currently has {values['duration'].unit} units."
+                f"currently has {value.unit} units."
             )
+        return value
diff --git a/stellarphot/settings/tests/test_models.py b/stellarphot/settings/tests/test_models.py
index c97fe235..3cb7ab6c 100644
--- a/stellarphot/settings/tests/test_models.py
+++ b/stellarphot/settings/tests/test_models.py
@@ -15,8 +15,7 @@ def test_create_aperture_settings_correctly():
     assert ap_set.radius == DEFAULT_APERTURE_SETTINGS["radius"]
     assert (
         ap_set.inner_annulus
-        == DEFAULT_APERTURE_SETTINGS["radius"] +
-        DEFAULT_APERTURE_SETTINGS["gap"]
+        == DEFAULT_APERTURE_SETTINGS["radius"] + DEFAULT_APERTURE_SETTINGS["gap"]
     )
     assert (
         ap_set.outer_annulus
@@ -38,10 +37,9 @@ def test_create_invalid_values(bad_one):
 DEFAULT_EXOPLANET_SETTINGS = dict(
     epoch=Time(0, format="jd"),
     period=0 * u.min,
-    identifier="",
+    identifier="a planet",
     coordinate=SkyCoord(
-        ra="00:00:00.00", dec="+00:00:00.0", frame="icrs",
-        unit=("hour", "degree")
+        ra="00:00:00.00", dec="+00:00:00.0", frame="icrs", unit=("hour", "degree")
     ),
     depth=0,
     duration=0 * u.min,
@@ -50,6 +48,7 @@ def test_create_invalid_values(bad_one):
 
 def test_create_exoplanet_correctly():
     planet = Exoplanet(**DEFAULT_EXOPLANET_SETTINGS)
+    print(planet)
     assert planet.epoch == DEFAULT_EXOPLANET_SETTINGS["epoch"]
     assert u.get_physical_type(planet.period) == "time"
     assert planet.identifier == DEFAULT_EXOPLANET_SETTINGS["identifier"]
@@ -59,17 +58,10 @@ def test_create_exoplanet_correctly():
 
 
 def test_create_invalid_exoplanet():
-    epoch = 1.0 * u.cm
-    period = 2.3
-    identifier = 'Bad planet'
-    coordinate = 'Not a real place'
-    depth = 23123
-    duration = 23213
+    values = DEFAULT_EXOPLANET_SETTINGS.copy()
+    # Make pediod and duration have invalid units for a time
+    values["period"] = values["period"].value * u.m
+    values["duration"] = values["duration"].value * u.m
     # Check that individual values that are bad raise an error
-    with pytest.raises(ValidationError, match='4 validation errors'):
-        Exoplanet(epoch=epoch,
-                  period=period,
-                  identifier=identifier,
-                  coordinate=coordinate,
-                  depth=depth,
-                  duration=duration)
+    with pytest.raises(ValidationError, match="2 validation errors"):
+        Exoplanet(**values)