mirror of
https://github.com/microsoft/regorus.git
synced 2026-08-05 02:16:11 +00:00
104 lines
2.7 KiB
Rust
104 lines
2.7 KiB
Rust
// Copyright (c) Microsoft Corporation.
|
|
// Licensed under the MIT License.
|
|
|
|
use anyhow::{bail, Result};
|
|
use regorus::*;
|
|
|
|
#[test]
|
|
fn extension() -> Result<()> {
|
|
fn repeat(mut params: Vec<Value>) -> Result<Value> {
|
|
match params.remove(0) {
|
|
Value::String(s) => {
|
|
let s = s.as_ref().to_owned();
|
|
Ok(Value::from(s.clone() + &s))
|
|
}
|
|
_ => bail!("param must be string"),
|
|
}
|
|
}
|
|
let mut engine = Engine::new();
|
|
engine.add_policy(
|
|
"test.rego".to_string(),
|
|
r#"package test
|
|
x = repeat("hello")
|
|
"#
|
|
.to_string(),
|
|
)?;
|
|
|
|
// Raises error since repeat is not defined.
|
|
assert!(engine.eval_query("data.test.x".to_string(), false).is_err());
|
|
|
|
// Register extension.
|
|
engine.add_extension("repeat".to_string(), 1, Box::new(repeat))?;
|
|
|
|
// Adding extension twice is error.
|
|
assert!(engine
|
|
.add_extension(
|
|
"repeat".to_string(),
|
|
1,
|
|
Box::new(|_| { Ok(Value::Undefined) })
|
|
)
|
|
.is_err());
|
|
|
|
let r = engine.eval_query("data.test.x".to_string(), false)?;
|
|
assert_eq!(
|
|
r.result[0].expressions[0].value.as_string()?.as_ref(),
|
|
"hellohello"
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
#[test]
|
|
fn extension_with_state() -> Result<()> {
|
|
#[derive(Clone)]
|
|
struct Gen {
|
|
n: i64,
|
|
}
|
|
|
|
let mut engine = Engine::new();
|
|
engine.add_policy(
|
|
"test.rego".to_string(),
|
|
r#"package test
|
|
x = gen()
|
|
"#
|
|
.to_string(),
|
|
)?;
|
|
|
|
let mut g = Box::new(Gen { n: 5 });
|
|
engine.add_extension(
|
|
"gen".to_string(),
|
|
0,
|
|
Box::new(move |_: Vec<Value>| {
|
|
let v = Value::from(g.n);
|
|
g.n += 1;
|
|
Ok(v)
|
|
}),
|
|
)?;
|
|
|
|
// First eval.
|
|
let r = engine.eval_query("data.test.x".to_string(), false)?;
|
|
assert_eq!(r.result[0].expressions[0].value.as_i64()?, 5);
|
|
|
|
// Second eval will produce a new value since for each query, the
|
|
// internal evaluation state of the interpreter is cleared.
|
|
// This might change in the future.
|
|
let r = engine.eval_query("data.test.x".to_string(), false)?;
|
|
assert_eq!(r.result[0].expressions[0].value.as_i64()?, 6);
|
|
|
|
// Clone the engine.
|
|
// This should also clone the stateful extension.
|
|
let mut engine1 = engine.clone();
|
|
|
|
// Both the engines should produce the same value.
|
|
let r = engine.eval_query("data.test.x".to_string(), false)?;
|
|
let r1 = engine1.eval_query("data.test.x".to_string(), false)?;
|
|
assert_eq!(
|
|
r.result[0].expressions[0].value,
|
|
r1.result[0].expressions[0].value
|
|
);
|
|
|
|
assert_eq!(r.result[0].expressions[0].value.as_i64()?, 7);
|
|
|
|
Ok(())
|
|
}
|