tokrs 0.2.2

A CLI tool for counting tokens
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102

use super::*;
use std::fs;

fn temp_dir(name: &str) -> PathBuf {
    let dir = std::env::temp_dir().join(format!(
        "tokrs-load-{name}-{}-{}",
        std::process::id(),
        SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos()
    ));
    fs::create_dir_all(&dir).unwrap();
    dir
}

#[test]
fn test_discover_files_respects_extension_and_depth() {
    let base = temp_dir("discover");
    fs::create_dir_all(base.join("a/b")).unwrap();
    fs::write(base.join("root.jsonl"), "{}").unwrap();
    fs::write(base.join("root.txt"), "{}").unwrap();
    fs::write(base.join("a/mid.jsonl"), "{}").unwrap();
    fs::write(base.join("a/b/deep.jsonl"), "{}").unwrap();

    let found = discover_files(&base, "jsonl", 1);
    let names: Vec<String> = found
        .iter()
        .map(|p| p.strip_prefix(&base).unwrap().to_string_lossy().to_string())
        .collect();
    assert_eq!(names, vec!["a/mid.jsonl", "root.jsonl"]);
    fs::remove_dir_all(&base).ok();
}

#[test]
fn test_discover_files_missing_base_is_empty() {
    let base = temp_dir("missing");
    fs::remove_dir_all(&base).ok();
    assert!(discover_files(&base, "jsonl", 3).is_empty());
}

#[test]
fn test_read_jsonl_skips_malformed_lines() {
    let path = temp_dir("jsonl").join("f.jsonl");
    fs::write(&path, "{\"a\":1}\nnot-json\n{\"a\":2}\n\n{\"a\":3}").unwrap();
    let rows = read_jsonl(&path).unwrap();
    assert_eq!(rows.len(), 3);
    assert_eq!(rows[2]["a"], 3);
}

#[test]
fn test_read_json_ok_and_corrupted() {
    let dir = temp_dir("readjson");
    let ok = dir.join("a.json");
    fs::write(&ok, "{\"x\":1}").unwrap();
    assert_eq!(read_json(&ok).unwrap()["x"], 1);
    let bad = dir.join("b.json");
    fs::write(&bad, "{oops").unwrap();
    assert!(matches!(read_json(&bad), Err(AppError::Corrupted { .. })));
}

#[test]
fn test_u64_and_str_get() {
    let v: Value = serde_json::from_str(r#"{"a":{"b":42},"s":{"t":"x"},"f":1.7}"#).unwrap();
    assert_eq!(u64_get(&v, &["a", "b"]), 42);
    assert_eq!(u64_get(&v, &["f"]), 1);
    assert_eq!(u64_get(&v, &["nope"]), 0);
    assert_eq!(str_get(&v, &["s", "t"]), Some("x"));
    assert_eq!(str_get(&v, &["a", "t"]), None);
}

#[test]
fn test_cost_get() {
    let v: Value =
        serde_json::from_str(r#"{"a":{"b":0.5},"z":{"b":0},"n":{"b":-1},"s":{"b":"x"}}"#).unwrap();
    assert_eq!(cost_get(&v, &["a", "b"]), Some(0.5));
    assert_eq!(cost_get(&v, &["z", "b"]), None);
    assert_eq!(cost_get(&v, &["n", "b"]), None);
    assert_eq!(cost_get(&v, &["s", "b"]), None);
    assert_eq!(cost_get(&v, &["nope"]), None);
}

#[test]
fn test_bool_get() {
    let v: Value = serde_json::from_str(r#"{"t":{"b":true},"x":{"b":1}}"#).unwrap();
    assert!(bool_get(&v, &["t", "b"]));
    assert!(!bool_get(&v, &["x", "b"]));
    assert!(!bool_get(&v, &["nope"]));
}

#[test]
fn test_timestamp_to_epoch() {
    let v: Value = serde_json::from_str(
        "[1767000000, 1767000000000, \"1767000000\", \"2026-09-01T00:00:00Z\"]",
    )
    .unwrap();
    let arr = v.as_array().unwrap();
    assert_eq!(timestamp_to_epoch(&arr[0]), Some(1_767_000_000));
    assert_eq!(timestamp_to_epoch(&arr[1]), Some(1_767_000_000));
    assert_eq!(timestamp_to_epoch(&arr[2]), Some(1_767_000_000));
    assert_eq!(timestamp_to_epoch(&arr[3]), Some(1_788_220_800));
}