1use anyhow::Result;
2use extism_pdk::{log, FromBytesOwned, LogLevel, Memory, ToMemory};
3
4mod harness {
5 #[link(wasm_import_module = "xtp:test/harness")]
6 extern "C" {
7 pub fn call(name: u64, input: u64) -> u64;
8 pub fn time(name: u64, input: u64) -> u64;
9 pub fn assert(name: u64, value: u64, message: u64);
10 pub fn reset();
11 pub fn group(name: u64);
12 pub fn mock_input() -> u64;
13 }
14}
15
16pub fn mock_input<T: FromBytesOwned>() -> Option<T> {
17 let offs = unsafe { harness::mock_input() };
18 let mem = Memory::find(offs)?;
19 let x = mem.to();
20 mem.free();
21 match x {
22 Ok(x) => Some(x),
23 Err(e) => {
24 log!(LogLevel::Error, "Invalid mock_input type: {:?}", e);
25 None
26 }
27 }
28}
29
30pub fn call_memory(func_name: impl AsRef<str>, input: impl ToMemory) -> Result<Memory> {
32 let func_name = func_name.as_ref();
33 let func_mem = Memory::from_bytes(func_name)?;
34 let input_mem = input.to_memory()?;
35 let output_ptr = unsafe { harness::call(func_mem.offset(), input_mem.offset()) };
36 func_mem.free();
37 input_mem.free();
38
39 let output = match Memory::find(output_ptr) {
40 None => anyhow::bail!("Error in call to {func_name}: invalid output offset"),
41 Some(x) => x,
42 };
43 Ok(output)
44}
45
46pub fn call<T: extism_pdk::FromBytesOwned>(
48 func_name: impl AsRef<str>,
49 input: impl ToMemory,
50) -> Result<T> {
51 let output_mem = call_memory(func_name, input)?;
52 let output = output_mem.to();
53 output_mem.free();
54 output
55}
56
57pub fn time_ns(func_name: impl AsRef<str>, input: impl ToMemory) -> Result<u64> {
59 let func_name = func_name.as_ref();
60 let func_mem = Memory::from_bytes(func_name)?;
61 let input_mem = input.to_memory()?;
62 let ns = unsafe { harness::time(func_mem.offset(), input_mem.offset()) };
63 func_mem.free();
64 input_mem.free();
65
66 Ok(ns)
67}
68
69pub fn time_sec(func_name: impl AsRef<str>, input: impl ToMemory) -> Result<f64> {
71 time_ns(func_name, input).map(|x| x as f64 / 1e9)
72}
73
74pub fn assert(name: impl AsRef<str>, outcome: bool, reason: impl AsRef<str>) {
77 let name_mem = Memory::from_bytes(name.as_ref()).expect("assert name Extism memory");
78 let reason_mem = Memory::from_bytes(reason.as_ref()).expect("assert reason Extism memory");
79 unsafe {
80 harness::assert(name_mem.offset(), outcome as u64, reason_mem.offset());
81 }
82 reason_mem.free();
83 name_mem.free();
84}
85
86pub fn assert_eq<U: std::fmt::Debug, T: std::fmt::Debug + PartialEq<U>>(
88 msg: impl AsRef<str>,
89 x: T,
90 y: U,
91) {
92 assert(msg, x == y, format!("Expected {:?} == {:?}", x, y));
93}
94
95pub fn assert_ne<U: std::fmt::Debug, T: std::fmt::Debug + PartialEq<U>>(
97 msg: impl AsRef<str>,
98 x: T,
99 y: U,
100) {
101 assert(msg, x != y, format!("Expected {:?} != {:?}", x, y));
102}
103
104pub fn assert_gt<U: std::fmt::Debug, T: std::fmt::Debug + PartialOrd<U>>(
106 msg: impl AsRef<str>,
107 x: T,
108 y: U,
109) {
110 assert(msg, x > y, format!("Expected {:?} > {:?}", x, y));
111}
112
113pub fn assert_gte<U: std::fmt::Debug, T: std::fmt::Debug + PartialOrd<U>>(
115 msg: impl AsRef<str>,
116 x: T,
117 y: U,
118) {
119 assert(msg, x >= y, format!("Expected {:?} >= {:?}", x, y));
120}
121
122pub fn assert_lt<U: std::fmt::Debug, T: std::fmt::Debug + PartialOrd<U>>(
124 msg: impl AsRef<str>,
125 x: T,
126 y: U,
127) {
128 assert(msg, x < y, format!("Expected {:?} < {:?}", x, y));
129}
130
131pub fn assert_lte<U: std::fmt::Debug, T: std::fmt::Debug + PartialOrd<U>>(
133 msg: impl AsRef<str>,
134 x: T,
135 y: U,
136) {
137 assert(msg, x <= y, format!("Expected {:?} <= {:?}", x, y));
138}
139
140fn start_group(name: impl AsRef<str>) {
142 let name_mem = Memory::from_bytes(name.as_ref()).expect("assert message Extism memory");
143 unsafe {
144 harness::group(name_mem.offset());
145 }
146 name_mem.free();
147}
148
149pub fn reset() {
151 unsafe {
152 harness::reset();
153 }
154}
155
156pub fn group(name: impl AsRef<str>, f: impl FnOnce() -> Result<()>) -> Result<()> {
169 reset();
170 start_group(name);
171 let res = f();
172 reset();
173 res
174}
175
176#[macro_export]
177macro_rules! assert {
178 ($name:expr, $b:expr) => {
179 $crate::assert(
180 $name,
181 $b,
182 format!("Assertion failed ({}:{})", file!(), line!()),
183 );
184 };
185}
186
187#[macro_export]
188macro_rules! assert_eq {
189 ($name:expr, $a:expr, $b:expr) => {
190 $crate::assert(
191 $name,
192 $a == $b,
193 format!("Expected {:?} == {:?} ({}:{})", $a, $b, file!(), line!()),
194 );
195 };
196}
197
198#[macro_export]
199macro_rules! assert_ne {
200 ($name:expr, $a:expr, $b:expr) => {
201 $crate::assert(
202 $name,
203 $a != $b,
204 format!("Expected {:?} != {:?} ({}:{})", $a, $b, file!(), line!()),
205 );
206 };
207}
208
209#[macro_export]
210macro_rules! assert_lt {
211 ($name:expr, $a:expr, $b:expr) => {
212 $crate::assert(
213 $name,
214 $a < $b,
215 format!("Expected {:?} < {:?} ({}:{})", $a, $b, file!(), line!()),
216 );
217 };
218}
219
220#[macro_export]
221macro_rules! assert_lte {
222 ($name:expr, $a:expr, $b:expr) => {
223 $crate::assert(
224 $name,
225 $a <= $b,
226 format!("Expected {:?} <= {:?} ({}:{})", $a, $b, file!(), line!()),
227 );
228 };
229}
230
231#[macro_export]
232macro_rules! assert_gt {
233 ($name:expr, $a:expr, $b:expr) => {
234 $crate::assert(
235 $name,
236 $a > $b,
237 format!("Expected {:?} > {:?} ({}:{})", $a, $b, file!(), line!()),
238 );
239 };
240}
241
242#[macro_export]
243macro_rules! assert_gte {
244 ($name:expr, $a:expr, $b:expr) => {
245 $crate::assert(
246 $name,
247 $a >= $b,
248 format!("Expected {:?} >= {:?} ({}:{})", $a, $b, file!(), line!()),
249 );
250 };
251}