|
18 | 18 | ability to be created from and serialized to YAML format. |
19 | 19 | """ |
20 | 20 |
|
| 21 | +import textwrap |
| 22 | + |
21 | 23 | from absl.testing import absltest |
22 | 24 | from cel_expr_python import cel |
| 25 | +from cel_expr_python.ext import ext_math |
23 | 26 | from cel.expr.conformance.proto2 import test_all_types_pb2 as test_all_types_pb |
24 | 27 |
|
25 | 28 |
|
@@ -91,6 +94,18 @@ def test_invalid_yaml(self): |
91 | 94 | str(e.exception), |
92 | 95 | ) |
93 | 96 |
|
| 97 | + def test_config_export_container(self): |
| 98 | + env = cel.NewEnv( |
| 99 | + container="test.container" |
| 100 | + ) |
| 101 | + yaml = env.config().to_yaml() |
| 102 | + self.assertEqual( |
| 103 | + normalize_yaml(yaml), |
| 104 | + normalize_yaml(""" |
| 105 | + container: "test.container" |
| 106 | + """), |
| 107 | + ) |
| 108 | + |
94 | 109 | def test_config_export_variables(self): |
95 | 110 | config = cel.NewEnv( |
96 | 111 | variables={ |
@@ -236,27 +251,78 @@ def test_config_variable_types(self): |
236 | 251 | self.assertEqual(res.type(), cel.Type.INT) |
237 | 252 | self.assertEqual(res.value(), 42) |
238 | 253 |
|
| 254 | + def test_config_extensions(self): |
| 255 | + config = cel.NewEnvConfigFromYaml(""" |
| 256 | + extensions: |
| 257 | + - name: math |
| 258 | + - name: strings |
| 259 | + """) |
| 260 | + env = cel.NewEnv( |
| 261 | + config=config, |
| 262 | + extensions=[TestCelExtension()], |
| 263 | + ) |
| 264 | + yaml = env.config().to_yaml() |
| 265 | + self.assertEqual( |
| 266 | + normalize_yaml(yaml), |
| 267 | + normalize_yaml(""" |
| 268 | + extensions: |
| 269 | + - name: "math" |
| 270 | + - name: "strings" |
| 271 | + - name: "test_cel_extension" |
| 272 | + """), |
| 273 | + ) |
| 274 | + res = env.compile("'%.4f'.format([math.sqrt(2)])").eval() |
| 275 | + self.assertEqual(res.value(), "1.4142") |
| 276 | + res = env.compile("hello('World')").eval() |
| 277 | + self.assertEqual(res.value(), "Hello, World!") |
| 278 | + |
| 279 | + def test_config_extensions_override(self): |
| 280 | + # TODO(b/498655870): add assertion based on extension aliases once |
| 281 | + # supported. |
| 282 | + config = cel.NewEnvConfigFromYaml(""" |
| 283 | + extensions: |
| 284 | + - name: cel.lib.ext.math |
| 285 | + version: 0 |
| 286 | + - name: cel.lib.ext.strings |
| 287 | + """) |
| 288 | + with self.assertRaises(Exception) as e: |
| 289 | + cel.NewEnv( |
| 290 | + config=config, |
| 291 | + extensions=[ext_math.ExtMath()], |
| 292 | + ) |
| 293 | + self.assertIn( |
| 294 | + "Extension 'cel.lib.ext.math' version 0 is already included. Cannot" |
| 295 | + " also include version 'latest'", |
| 296 | + str(e.exception), |
| 297 | + ) |
| 298 | + |
| 299 | + |
| 300 | +class TestCelExtension(cel.CelExtension): |
| 301 | + """An example CEL extension for testing.""" |
| 302 | + |
| 303 | + def __init__(self): |
| 304 | + super().__init__( |
| 305 | + "test_cel_extension", |
| 306 | + functions=[ |
| 307 | + cel.FunctionDecl( |
| 308 | + "hello", |
| 309 | + [ |
| 310 | + cel.Overload( |
| 311 | + "hello(string)", |
| 312 | + return_type=cel.Type.STRING, |
| 313 | + parameters=[ |
| 314 | + cel.Type.STRING, |
| 315 | + ], |
| 316 | + impl=lambda arg: f"Hello, {arg}!", |
| 317 | + ) |
| 318 | + ], |
| 319 | + ), |
| 320 | + ], |
| 321 | + ) |
| 322 | + |
239 | 323 |
|
240 | 324 | def normalize_yaml(yaml: str) -> str: |
241 | | - lines = yaml.split("\n") |
242 | | - indent = -1 |
243 | | - unindented_lines = [] |
244 | | - for line in lines: |
245 | | - pos = -1 |
246 | | - for i, char in enumerate(line): |
247 | | - if char != " " and char != "\t": |
248 | | - pos = i |
249 | | - break |
250 | | - if pos == -1: |
251 | | - # Skip blank lines. |
252 | | - continue |
253 | | - if indent == -1: |
254 | | - indent = pos |
255 | | - if pos >= indent: |
256 | | - unindented_lines.append(line[indent:]) |
257 | | - else: |
258 | | - unindented_lines.append(line) |
259 | | - return "\n".join(unindented_lines) |
| 325 | + return textwrap.dedent(yaml).strip() |
260 | 326 |
|
261 | 327 |
|
262 | 328 | if __name__ == "__main__": |
|
0 commit comments