ta_lib_in_rust/indicators/
test_util.rs

1use polars::prelude::*;
2
3/// Creates a test OHLCV DataFrame for testing indicator functions
4///
5/// This function generates a DataFrame with OHLCV data suitable for testing technical indicators.
6/// The data follows a simple pattern with some randomness to simulate realistic price movements.
7///
8/// # Returns
9///
10/// * `DataFrame` - A DataFrame with columns "open", "high", "low", "close", "volume"
11pub fn create_test_ohlcv_df() -> DataFrame {
12    // Generate some sample price data
13    let base_price = 100.0;
14    let mut prices = Vec::new();
15    let mut open = Vec::new();
16    let mut high = Vec::new();
17    let mut low = Vec::new();
18    let mut close = Vec::new();
19    let mut volume = Vec::new();
20
21    for i in 0..100 {
22        // Create a simple sine wave pattern with some noise
23        let time_factor = i as f64 * 0.1;
24        let wave = (time_factor.sin() * 10.0) + base_price;
25        let noise = (i % 7) as f64 * 0.5;
26
27        let price = wave + noise;
28        prices.push(price);
29
30        // Generate OHLC based on the price
31        let o = price - 1.0 + (i % 5) as f64 * 0.2;
32        let c = price + 0.5 - (i % 3) as f64 * 0.3;
33        let h = price.max(o).max(c) + 1.0 + (i % 4) as f64 * 0.4;
34        let l = price.min(o).min(c) - 1.0 - (i % 6) as f64 * 0.2;
35
36        open.push(o);
37        high.push(h);
38        low.push(l);
39        close.push(c);
40
41        // Generate volume
42        let v = 1000.0 + (i % 10) as f64 * 200.0;
43        volume.push(v);
44    }
45
46    // Create the DataFrame
47
48    DataFrame::new(vec![
49        Series::new("open".into(), open).into(),
50        Series::new("high".into(), high).into(),
51        Series::new("low".into(), low).into(),
52        Series::new("close".into(), close).into(),
53        Series::new("volume".into(), volume).into(),
54    ])
55    .unwrap()
56}