diff --git a/straxen/config/protocols.py b/straxen/config/protocols.py index 76d6773fe..b351da166 100644 --- a/straxen/config/protocols.py +++ b/straxen/config/protocols.py @@ -65,6 +65,8 @@ def read_json(content: str, **kwargs): @URLConfig.register("take") def get_key(container: Container, take=None, **kwargs): """Return a single element of a container.""" + if not isinstance(container, dict): + raise ValueError(f"Container is not a dict but a {type(container)}") if take is None: return container if not isinstance(take, list): diff --git a/tests/test_url_config.py b/tests/test_url_config.py index 12af3476f..d004b4f61 100644 --- a/tests/test_url_config.py +++ b/tests/test_url_config.py @@ -25,6 +25,12 @@ def generate_random(_): return random.random() +@straxen.URLConfig.register("range") +def generate_range(length): + length = int(length) + return np.arange(length) + + @straxen.URLConfig.register("unpicklable") def return_lamba(_): return lambda x: x @@ -118,9 +124,9 @@ def test_cmt_protocol(self): self.assertTrue(abs(p.test_config - 219203.49884000001) < 1e-2) def test_json_protocol(self): - self.st.set_config({"test_config": "json://[1,2,3]"}) + self.st.set_config({"test_config": 'json://{"a":0}'}) p = self.st.get_single_plugin(nt_test_run_id, "test_data") - self.assertEqual(p.test_config, [1, 2, 3]) + self.assertEqual(p.test_config, {"a": 0}) def test_format_protocol(self): self.st.set_config({"test_config": "format://{run_id}?run_id=plugin.run_id"}) @@ -136,15 +142,16 @@ def test_fsspec_protocol(self): p = self.st.get_single_plugin(nt_test_run_id, "test_data") self.assertEqual(p.test_config, 999) - def test_chained(self): - self.st.set_config({"test_config": "take://json://[1,2,3]?take=0"}) - p = self.st.get_single_plugin(nt_test_run_id, "test_data") - self.assertEqual(p.test_config, 1) - def test_take_nested(self): - self.st.set_config({"test_config": 'take://json://{"a":[1,2,3]}?take=a&take=0'}) + self.st.set_config( + { + "test_config": ( + 'take://json://{"a":{"aa":0,"ab":1},"b":{"ba":2,"bb":3}}?take=b&take=ba' + ) + } + ) p = self.st.get_single_plugin(nt_test_run_id, "test_data") - self.assertEqual(p.test_config, 1) + self.assertEqual(p.test_config, 2) @unittest.skipIf(not straxen.utilix_is_configured(), "No db access, cannot test!") def test_bodedga_get(self): @@ -434,9 +441,9 @@ def test_regex_url_warnings(self): def test_pad_array(self): """Test that pad_array works as expected.""" - + n = 3 self.st.set_config( - {"test_config": "pad-array://json://[1,2,3]?pad_left=2&pad_right=3&pad_value=0"} + {"test_config": f"pad-array://range://{n}?pad_left=2&pad_right=3&pad_value=0"} ) p = self.st.get_single_plugin(nt_test_run_id, "test_data") self.assertEqual(len(p.test_config), 8)