From b85fa59c7bbd9710649115a259f8f94f5b8b5219 Mon Sep 17 00:00:00 2001
From: Metin San <metinisan@gmail.com>
Date: Sun, 28 Apr 2024 10:34:45 +0200
Subject: [PATCH] Add tests to get back to 99% coverage

---
 .pre-commit-config.yaml  |  1 +
 tests/test_validators.py | 25 +++++++++++++++++++++++++
 zodipy/_bandpass.py      | 13 ++++++-------
 zodipy/_ipd_model.py     |  5 +++--
 zodipy/_types.py         | 13 -------------
 5 files changed, 35 insertions(+), 22 deletions(-)
 delete mode 100644 zodipy/_types.py

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index dfed5f2..a2549ed 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -16,6 +16,7 @@ repos:
     rev: v1.10.0
     hooks:
     -   id: mypy
+        exclude: tests/
         additional_dependencies: [numpy>=1.21.0]
 hooks:
   - id: poetry-export
diff --git a/tests/test_validators.py b/tests/test_validators.py
index f4e5832..d429233 100644
--- a/tests/test_validators.py
+++ b/tests/test_validators.py
@@ -136,3 +136,28 @@ def test_interp_kind() -> None:
         model.get_emission_pix(
             pixels=[1, 4, 5], freq=27 * units.micron, nside=32, obs_time=OBS_TIME
         )
+
+
+def test_wrong_frame() -> None:
+    """Tests that an error is correctly raised when an incorrect frame is passed in."""
+    model = Zodipy()
+    with pytest.raises(ValueError):
+        model.get_emission_pix(
+            [1, 4, 5],
+            nside=32,
+            freq=27 * units.micron,
+            obs_time=OBS_TIME,
+            coord_in="not a valid frame",
+        )
+
+
+def test_non_quantity_ang_raises_error() -> None:
+    """Tests that an error is correctly raised if the user inputed angles are not Quantities."""
+    model = Zodipy()
+    with pytest.raises(TypeError):
+        model.get_emission_ang(
+            32,
+            12,
+            freq=27 * units.micron,
+            obs_time=OBS_TIME,
+        )
diff --git a/zodipy/_bandpass.py b/zodipy/_bandpass.py
index e5ecc5e..fba3428 100644
--- a/zodipy/_bandpass.py
+++ b/zodipy/_bandpass.py
@@ -3,8 +3,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
-import astropy.units as u
 import numpy as np
+from astropy import units
 
 from zodipy._constants import (
     MAX_INTERPOLATION_GRID_TEMPERATURE,
@@ -18,12 +18,11 @@
     import numpy.typing as npt
 
     from zodipy._ipd_model import InterplanetaryDustModel
-    from zodipy._types import FrequencyOrWavelength
 
 
 @dataclass
 class Bandpass:
-    frequencies: FrequencyOrWavelength
+    frequencies: units.Quantity
     weights: npt.NDArray[np.float64]
 
     def integrate(self, quantity: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
@@ -33,8 +32,8 @@ def integrate(self, quantity: npt.NDArray[np.float64]) -> npt.NDArray[np.float64
     def switch_convention(self) -> None:
         """Switch the bandpass from frequency to wavelength or the other way around."""
         self.frequencies = self.frequencies.to(
-            u.micron if self.frequencies.unit.is_equivalent(u.Hz) else u.Hz,
-            equivalencies=u.spectral(),
+            units.micron if self.frequencies.unit.is_equivalent(units.Hz) else units.Hz,
+            equivalencies=units.spectral(),
         )
         if self.frequencies.size > 1:
             self.frequencies = np.flip(self.frequencies)
@@ -43,7 +42,7 @@ def switch_convention(self) -> None:
 
 
 def validate_and_get_bandpass(
-    freq: FrequencyOrWavelength,
+    freq: units.Quantity,
     weights: npt.ArrayLike | None,
     model: InterplanetaryDustModel,
     extrapolate: bool,
@@ -62,7 +61,7 @@ def get_bandpass_interpolation_table(
     max_temp: float = MAX_INTERPOLATION_GRID_TEMPERATURE,
 ) -> npt.NDArray[np.float64]:
     """Pre-compute the bandpass integrated blackbody emission for a grid of temperatures."""
-    if not bandpass.frequencies.unit.is_equivalent(u.Hz):
+    if not bandpass.frequencies.unit.is_equivalent(units.Hz):
         bandpass.switch_convention()
 
     bandpass_integrals = np.zeros(n_points)
diff --git a/zodipy/_ipd_model.py b/zodipy/_ipd_model.py
index 4589106..f14fd47 100644
--- a/zodipy/_ipd_model.py
+++ b/zodipy/_ipd_model.py
@@ -5,8 +5,9 @@
 from typing import TYPE_CHECKING, Mapping, Sequence
 
 if TYPE_CHECKING:
+    from astropy import units
+
     from zodipy._ipd_comps import Component, ComponentLabel
-    from zodipy._types import FrequencyOrWavelength
 
 
 @dataclass
@@ -20,7 +21,7 @@ class InterplanetaryDustModel(ABC):
     """
 
     comps: Mapping[ComponentLabel, Component]
-    spectrum: FrequencyOrWavelength
+    spectrum: units.Quantity
 
     def to_dict(self) -> dict:
         """Return a dictionary representation of the model."""
diff --git a/zodipy/_types.py b/zodipy/_types.py
deleted file mode 100644
index c887a5a..0000000
--- a/zodipy/_types.py
+++ /dev/null
@@ -1,13 +0,0 @@
-from __future__ import annotations
-
-from typing import Sequence, Union
-
-import astropy.units as u
-import numpy as np
-import numpy.typing as npt
-
-Pixels = Union[int, Sequence[int], npt.NDArray[np.integer]]
-SkyAngles = Union[u.Quantity[u.deg], u.Quantity[u.rad]]
-FrequencyOrWavelength = Union[u.Quantity[u.Hz], u.Quantity[u.m]]
-ParameterDict = dict
-Float = Union[float, np.floating]