break out wip signal tests into separate test instances

Author: sgsmob

_delphi_utils_python/tests/test_signal.py

@@ -1,38 +1,34 @@
 """Tests for delphi_utils.signal."""
-import pandas as pd
 from unittest.mock import patch
+import pandas as pd
 
 from delphi_utils.signal import add_prefix, public_signal
 
 # Constants for mocking out the call to `covidcast.metadata` within `public_signal()`.
-SIGNALS = ["sig1", "sig2", "sig3"]
-SIGNALS_FRAME = pd.DataFrame(data={"signal": SIGNALS})
+PUBLIC_SIGNALS = ["sig1", "sig2", "sig3"]
+PUBLIC_SIGNALS_FRAME = pd.DataFrame(data={"signal": PUBLIC_SIGNALS})
 
 class TestSignal:
     """Tests for signal.py."""
 
+    def test_add_prefix_to_all(self):
+        """Tests that `add_prefix()` derives work-in-progress names for all input signals."""
+        assert add_prefix(["sig1", "sig3"], True, prefix="wip_") == ["wip_sig1", "wip_sig3"]
+    
+    def test_add_prefix_to_specified(self):
+        """Tests that `add_prefix()` derives work-in-progress names for specified signals."""
+        assert add_prefix(["sig1", "sig2", "sig3"], ["sig2"], prefix="wip_") ==\
+            ["sig1", "wip_sig2", "sig3"]
+    
     @patch("covidcast.metadata")
-    def test_handle_wip_signal(self, metadata):
-        """Tests that `add_prefix()` derives work-in-progress signals."""
-        metadata.return_value = SIGNALS_FRAME
-
-        # Test wip_signal = True
-        signal_names = SIGNALS
-        signal_names = add_prefix(SIGNALS, True, prefix="wip_")
-        assert all(s.startswith("wip_") for s in signal_names)
-        # Test wip_signal = list
-        signal_names = add_prefix(SIGNALS, [SIGNALS[0]], prefix="wip_")
-        assert signal_names[0].startswith("wip_")
-        assert all(not s.startswith("wip_") for s in signal_names[1:])
-        # Test wip_signal = False
-        signal_names = add_prefix(["xyzzy", SIGNALS[0]], False, prefix="wip_")
-        assert signal_names[0].startswith("wip_")
-        assert all(not s.startswith("wip_") for s in signal_names[1:])
+    def test_add_prefix_to_non_public(self, metadata):
+        """Tests that `add_prefix()` derives work-in-progress names for non-public signals."""
+        metadata.return_value = PUBLIC_SIGNALS_FRAME
+        assert add_prefix(["sig0", "sig1"], False, prefix="wip_") == ["wip_sig0", "sig1"]
 
     @patch("covidcast.metadata")
     def test_public_signal(self, metadata):
         """Tests that `public_signal()` identifies public vs. private signals."""
-        metadata.return_value = SIGNALS_FRAME
-
+        metadata.return_value = PUBLIC_SIGNALS_FRAME
         assert not public_signal("sig0")
         assert public_signal("sig2")