Minor changes in process.py, updated test cases

Author: vishakha1812

safegraph/delphi_safegraph/process.py

@@ -26,20 +26,23 @@ def add_prefix(signal_names, wip_signal, prefix: str):
     List of signal names
         wip/non wip signals for further computation
     """
-    if wip_signal is not None:
-        if wip_signal and isinstance(wip_signal, bool):
-            return [
-                (prefix + signal) if public_signal(signal)
-                else signal
-                for signal in signal_names
-            ]
-        if isinstance(wip_signal, list):
-            for signal in wip_signal:
-                if public_signal(signal):
-                    new_list = [prefix + signal]
-                    signal_names.remove(signal)
-            signal_names.extend(new_list)
-    return signal_names
+
+    if wip_signal in ("", False):
+        return signal_names
+    elif wip_signal and isinstance(wip_signal, bool):
+        return [
+            (prefix + signal) if public_signal(signal)
+            else signal
+            for signal in signal_names
+        ]
+    elif isinstance(wip_signal, list):
+        for signal in wip_signal:
+            if public_signal(signal):
+                signal_names.append(prefix + signal)
+                signal_names.remove(signal)
+        return signal_names
+    else:
+        raise ValueError("Supply True | False or '' or [] | list()")
 
 
 # Check if the signal name is public

safegraph/tests/test_process.py

@@ -8,7 +8,7 @@
 from delphi_safegraph.process import (
     construct_signals,
     aggregate,
-    signal_name
+    add_prefix
 )
 from delphi_safegraph.run import SIGNALS
 from delphi_utils import read_params
@@ -49,11 +49,13 @@ def test_aggregate_state(self):
         x = df[f'{signal_names[0]}_se'].values
         assert np.all(x[~np.isnan(x)] >= 0)
 
-    def test_signal_name(self):
-        assert read_params()["wip_signal"] is not None, "supply value in params"
-        assert type(read_params()["wip_signal"]) == list or type(read_params()["wip_signal"]) == bool, "Supply True|False|list()"
-        signals = signal_name(signal_names, wip_signal=read_params()['wip_signal'],prefix='wip_')
-        assert (len(signals) >= len(signal_names))
+    def test_handle_wip_signal(self):
+        wip_signal = read_params()["wip_signal"]
+        assert isinstance(wip_signal, (list, bool)) or wip_signal == "", "Supply True | False or "" or [] | list()"
+        if isinstance(wip_signal, list):
+            assert set(wip_signal).issubset(set(SIGNALS)), "signal in params don't belong in the registry"
+        updated_signal_names = add_prefix(SIGNALS, wip_signal, prefix='wip_')
+        assert (len(updated_signal_names) >= len(SIGNALS))