Add tests for _get_validation_warnings

src/transport_performance/gtfs/gtfs_utils.py

@@ -320,7 +320,9 @@ def convert_pandas_to_plotly(
     return fig
 
 
-def _get_validation_warnings(gtfs, message: str) -> pd.DataFrame:
+def _get_validation_warnings(
+    gtfs, message: str, return_type: str = "values"
+) -> pd.DataFrame:
     """Get warnings from the validity_df table based on a regex.
 
     Parameters
@@ -329,6 +331,9 @@ def _get_validation_warnings(gtfs, message: str) -> pd.DataFrame:
         The gtfs instance to obtain the warnings from.
     message : str
         The regex to use for filtering the warnings.
+    return_type : str, optional
+        The return type of the warnings. Can be eithher 'values' or 'dataframe'
+         by default 'values'
 
     Returns
     -------
@@ -346,16 +351,18 @@ def _get_validation_warnings(gtfs, message: str) -> pd.DataFrame:
         ),
     )
     _type_defence(message, "message", str)
-    needed_warnings = (
-        gtfs.validity_df[
-            gtfs.validity_df["message"].str.contains(
-                message, regex=True, na=False
-            )
-        ]
-        .copy()
-        .values
-    )
-    return needed_warnings
+    return_type = return_type.lower().strip()
+    if return_type not in ["values", "dataframe"]:
+        raise ValueError(
+            "'return_type' expected one of ['values', 'dataframe]"
+            f". Got {return_type}"
+        )
+    needed_warnings = gtfs.validity_df[
+        gtfs.validity_df["message"].str.contains(message, regex=True, na=False)
+    ].copy()
+    if return_type == "dataframe":
+        return needed_warnings
+    return needed_warnings.values
 
 
 def _remove_validation_row(

tests/gtfs/test_gtfs_utils.py

@@ -8,13 +8,15 @@
 import geopandas as gpd
 from shapely.geometry import box
 from plotly.graph_objects import Figure as PlotlyFigure
+import numpy as np
 
 from transport_performance.gtfs.validation import GtfsInstance
 from transport_performance.gtfs.gtfs_utils import (
     bbox_filter_gtfs,
     _add_validation_row,
     filter_gtfs_around_trip,
     convert_pandas_to_plotly,
+    _get_validation_warnings,
 )
 
 # location of GTFS test fixture
@@ -203,3 +205,81 @@ def test_convert_pandas_to_plotly_on_pass(self, test_df):
             "Expected type plotly.graph_objects.Figure but "
             f"{type(fig_return)} found"
         )
+
+
+class TestGetValidationWarnings(object):
+    """Tests for _get_validation_warnings."""
+
+    def test__get_validation_warnings_defence(self):
+        """Test thhe defences of _get_validation_warnings."""
+        with pytest.raises(
+            TypeError, match=".* expected a GtfsInstance object"
+        ):
+            _get_validation_warnings(True, "test_msg")
+        gtfs = GtfsInstance(gtfs_pth=GTFS_FIX_PTH)
+        with pytest.raises(
+            AttributeError, match="The gtfs has not been validated.*"
+        ):
+            _get_validation_warnings(gtfs, "test")
+        gtfs.is_valid()
+        with pytest.raises(
+            ValueError, match=r"'return_type' expected one of \[.*\]\. Got .*"
+        ):
+            _get_validation_warnings(gtfs, "tester", "tester")
+
+    def test__get_validation_warnings(self):
+        """Test _get_validation_warnings on pass."""
+        gtfs = GtfsInstance(GTFS_FIX_PTH)
+        gtfs.is_valid()
+        # test return types
+        df_exp = _get_validation_warnings(
+            gtfs, "test", return_type="dataframe"
+        )
+        assert isinstance(
+            df_exp, pd.DataFrame
+        ), f"Expected df, got {type(df_exp)}"
+        ndarray_exp = _get_validation_warnings(gtfs, "test")
+        assert isinstance(
+            ndarray_exp, np.ndarray
+        ), f"Expected np.ndarray, got {type(ndarray_exp)}"
+        # test with valld regex (assertions on DF data without DF)
+        regex_matches = _get_validation_warnings(
+            gtfs, "Unrecognized column *.", return_type="dataframe"
+        )
+        assert len(regex_matches) == 5, (
+            "Getting validaiton warnings returned"
+            "unexpected number of warnings"
+        )
+        assert list(regex_matches["type"].unique()) == [
+            "warning"
+        ], "Dataframe type column not asd expected"
+        assert list(regex_matches.table) == [
+            "agency",
+            "stop_times",
+            "stops",
+            "trips",
+            "trips",
+        ], "Dataframe table column not as expected"
+        # test with matching message (no regex)
+        exact_match = _get_validation_warnings(
+            gtfs, "Unrecognized column agency_noc", return_type="Dataframe"
+        )
+        assert list(exact_match.values[0]) == [
+            "warning",
+            "Unrecognized column agency_noc",
+            "agency",
+            [],
+        ], "Dataframe values not as expected"
+        assert (
+            len(exact_match) == 1
+        ), f"Expected one match, found {len(exact_match)}"
+        # test with no matches (regex)
+        regex_no_match = _get_validation_warnings(
+            gtfs, ".*This is a test.*", return_type="Dataframe"
+        )
+        assert len(regex_no_match) == 0, "No matches expected. Matches found"
+        # test with no match (no regex)
+        no_match = _get_validation_warnings(
+            gtfs, "This is a test!!!", return_type="Dataframe"
+        )
+        assert len(no_match) == 0, "No matches expected. Matched found"