Repository: peopledoc/workalendar
Pull request #624 (2021-02-05 13:49:32)

Improve coverage for the astronomy.py module
refs #546

- [x] Tests covering all corner cases
- [x] Changelog amended with a mention describing your changes.


Patch:
diff --git a/Changelog.md b/Changelog.md
index 7368695..6de328b 100644
--- a/Changelog.md
+++ b/Changelog.md
@@ -8,6 +8,7 @@
 - Improve Russia coverage (#546).
 - Improve USA calendar coverage by removing a method that wasn't used anyways (`get_washington_birthday_december()`). The method is implemented in both Indiana and Georgia State calendars, and is specific for each state, even if they look very similar (#546).
 - Temporary downgrade of `pyupgrade` to fix the `pyup_dirs`.
+- Improve the `astronomy.py` module coverage (#546).
 
 ## v14.3.0 (2021-01-15)
 
diff --git a/workalendar/astronomy.py b/workalendar/astronomy.py
index ba3c53f..7617ecc 100644
--- a/workalendar/astronomy.py
+++ b/workalendar/astronomy.py
@@ -42,21 +42,39 @@ def get_current_longitude(current_date, earth, sun):
     return longitude.radians
 
 
-def newton(f, x0, x1, precision=newton_precision):
+def newton(f, x0, x1, precision=newton_precision,
+           **func_kwargs):
     """Return an x-value at which the given function reaches zero.
 
     Stops and declares victory once the x-value is within ``precision``
     of the solution, which defaults to a half-second of clock time.
-
     """
-    f0, f1 = f(x0), f(x1)
+    f0, f1 = f(x0, **func_kwargs), f(x1, **func_kwargs)
     while f1 and abs(x1 - x0) > precision and f1 != f0:
         new_x1 = x1 + (x1 - x0) / (f0 / f1 - 1)
         x0, x1 = x1, new_x1
-        f0, f1 = f1, f(x1)
+        f0, f1 = f1, f(x1, **func_kwargs)
     return x1
 
 
+def newton_angle_function(t, ts, target_angle, body1, body2):
+    """
+    Compute the longitude of body2 relative to body1
+
+    In our case, it's Earth & Sun, but it could be used as any other
+    combination of solar system planets/bodies.
+    """
+    # We've got a float which is the `tt`
+    sky_tt = ts.tt_jd(t)
+    longitude = get_current_longitude(sky_tt, body1, body2)
+    result = target_angle - longitude
+    if result > pi:
+        result = result - pi
+    if result < -pi:
+        result = result + pi
+    return result
+
+
 def solar_term(year, degrees, timezone='UTC'):
     """
     Returns the date of the solar term for the given longitude
@@ -92,17 +110,6 @@ def solar_term(year, degrees, timezone='UTC'):
     # convert to "tt" and reconvert it back to a Time object
     t0 = ts.tt_jd(jan_first.tt + date_delta)
 
-    def f(t):
-        # We've got a float which is the `tt`
-        sky_tt = ts.tt_jd(t)
-        longitude = get_current_longitude(sky_tt, earth, sun)
-        result = target_angle - longitude
-        if result > pi:
-            result = result - pi
-        elif result < -pi:
-            result = result + pi
-        return result
-
     # Using datetimes to compute the next step date
     t0_plus_one_minute = t0.utc_datetime() + timedelta(minutes=1)
     # Back to Skyfield Time objects
@@ -113,7 +120,15 @@ def solar_term(year, degrees, timezone='UTC'):
     # Adding one minute to have a second boundary
     t0_plus_one_minute = t0_plus_one_minute.tt
     # Newton method to converge towards the target angle
-    t = newton(f, t0, t0_plus_one_minute)
+    t = newton(
+        newton_angle_function, t0, t0_plus_one_minute,
+        precision=newton_precision,
+        # Passed as kwargs to the angle function
+        ts=ts,
+        target_angle=target_angle,
+        body1=earth,
+        body2=sun,
+    )
     # Here we have a float to convert to julian days.
     t = ts.tt_jd(t)
     # To convert to datetime


Test patch:
diff --git a/workalendar/tests/test_astronomy.py b/workalendar/tests/test_astronomy.py
index 5a57251..52fda61 100644
--- a/workalendar/tests/test_astronomy.py
+++ b/workalendar/tests/test_astronomy.py
@@ -1,6 +1,12 @@
+import pytest
 from datetime import date
+from math import pi
+from unittest.mock import patch
 
-from ..astronomy import calculate_equinoxes, solar_term
+from skyfield.api import Loader
+from skyfield_data import get_skyfield_data_path
+
+from ..astronomy import calculate_equinoxes, solar_term, newton_angle_function
 
 
 def test_calculate_some_equinoxes():
@@ -27,3 +33,62 @@ def test_qingming_festivals_hk():
     assert solar_term(2019, 15, 'Asia/Hong_Kong') == date(2019, 4, 5)
     assert solar_term(2020, 15, 'Asia/Hong_Kong') == date(2020, 4, 4)
     assert solar_term(2021, 15, 'Asia/Hong_Kong') == date(2021, 4, 4)
+
+
+@pytest.fixture(scope='session')
+def params_newton_angle():
+    """
+    Session-scoped fixture to "cache" the newton angle func parameters
+    """
+    load = Loader(get_skyfield_data_path())
+    ts = load.timescale()
+    planets = load('de421.bsp')
+    earth = planets['earth']
+    sun = planets['sun']
+
+    jan_first = ts.utc(date(2021, 1, 1))
+    t0 = ts.tt_jd(jan_first.tt).tt
+    return t0, ts, earth, sun
+
+
+def test_newton_angle_function_normal_range(params_newton_angle):
+    """
+    Test the newton angle func when the longitude is in the range [-pi, +pi].
+    """
+    t0, ts, earth, sun = params_newton_angle
+
+    with patch('workalendar.astronomy.get_current_longitude') as patched:
+        patched.return_value = pi
+        assert newton_angle_function(t0, ts, 0, earth, sun) == -pi
+
+
+def test_newton_angle_function_above_pi(params_newton_angle):
+    """
+    Test the newton angle function when the longitude is > +pi.
+
+    This should not happen, but it was implemented in the function, so in order
+    to make sure that the resulting angle is always in the range [-pi, +pi],
+    we added these tests with those out of range values.
+    """
+    t0, ts, earth, sun = params_newton_angle
+
+    with patch('workalendar.astronomy.get_current_longitude') as patched:
+        patched.return_value = pi + 0.1
+        expected = pytest.approx(-.1)
+        assert newton_angle_function(t0, ts, 0, earth, sun) == expected
+
+
+def test_newton_angle_function_below_minus_pi(params_newton_angle):
+    """
+    Test the newton angle function when the longitude is < -pi.
+
+    This should not happen, but it was implemented in the function, so in order
+    to make sure that the resulting angle is always in the range [-pi, +pi],
+    we added these tests with those out of range values.
+    """
+    t0, ts, earth, sun = params_newton_angle
+
+    with patch('workalendar.astronomy.get_current_longitude') as patched:
+        patched.return_value = -pi - 0.1
+        expected = pytest.approx(.1)
+        assert newton_angle_function(t0, ts, 0, earth, sun) == expected