1c7194176a
Add compressFileZstd, decompressZstd, and walkZstFiles helpers using klauspost/compress. Promote zstd from indirect to direct dep. Co-Authored-By: Virgil <virgil@lethean.io>
58 lines
1.4 KiB
Go
58 lines
1.4 KiB
Go
package lem
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
)
|
|
|
|
func TestDecompressZstd(t *testing.T) {
|
|
// Create a temp .jsonl file and compress it
|
|
dir := t.TempDir()
|
|
src := filepath.Join(dir, "test.jsonl")
|
|
os.WriteFile(src, []byte(`{"messages":[{"role":"user","content":"hello"}]}`+"\n"), 0644)
|
|
|
|
// Compress with Go zstd
|
|
zstPath := src + ".zst"
|
|
if err := compressFileZstd(src, zstPath); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
os.Remove(src) // remove original
|
|
|
|
// Decompress
|
|
outPath := filepath.Join(dir, "out.jsonl")
|
|
if err := decompressZstd(zstPath, outPath); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
data, err := os.ReadFile(outPath)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(data) != `{"messages":[{"role":"user","content":"hello"}]}`+"\n" {
|
|
t.Fatalf("unexpected content: %q", string(data))
|
|
}
|
|
}
|
|
|
|
func TestWalkZstFiles(t *testing.T) {
|
|
dir := t.TempDir()
|
|
// Create a nested .jsonl.zst
|
|
sub := filepath.Join(dir, "sub")
|
|
os.MkdirAll(sub, 0755)
|
|
src := filepath.Join(sub, "data.jsonl")
|
|
os.WriteFile(src, []byte("line1\nline2\n"), 0644)
|
|
compressFileZstd(src, src+".zst")
|
|
os.Remove(src)
|
|
|
|
// Also create a .json file (should be ignored)
|
|
os.WriteFile(filepath.Join(sub, "probes.json"), []byte("[]"), 0644)
|
|
|
|
var found []string
|
|
walkZstFiles(dir, func(zstPath string) error {
|
|
found = append(found, zstPath)
|
|
return nil
|
|
})
|
|
if len(found) != 1 {
|
|
t.Fatalf("expected 1 .zst file, got %d: %v", len(found), found)
|
|
}
|
|
}
|