fastx-io 0.3.0

//! End-to-end tests that go through real files, including gzip.

use std::fs;
use std::path::{Path, PathBuf};

use fastx::{Alphabet, Error, FastaIndex, Format, IndexedFasta, Sequence, WriterBuilder};

/// A per-test scratch directory that is removed on drop.
struct Scratch(PathBuf);

impl Scratch {
    fn new(tag: &str) -> Scratch {
        let dir = std::env::temp_dir().join(format!("fastx-it-{}-{tag}", std::process::id()));
        let _ = fs::remove_dir_all(&dir);
        fs::create_dir_all(&dir).expect("create scratch dir");
        Scratch(dir)
    }

    fn path(&self, name: &str) -> PathBuf {
        self.0.join(name)
    }
}

impl Drop for Scratch {
    fn drop(&mut self) {
        let _ = fs::remove_dir_all(&self.0);
    }
}

fn corpus(records: usize) -> Vec<Sequence> {
    (0..records)
        .map(|i| {
            let seq: Vec<u8> = (0..(i % 97) + 1).map(|j| b"ACGTN"[(i + j) % 5]).collect();
            let quality: Vec<u8> = seq
                .iter()
                .enumerate()
                .map(|(j, _)| b'!' + (j % 40) as u8)
                .collect();
            Sequence::fastq(format!("read{i}"), seq, quality)
                .unwrap()
                .with_description(format!("length={} index={i}", (i % 97) + 1))
        })
        .collect()
}

#[test]
fn fastq_file_round_trip() {
    let scratch = Scratch::new("fq");
    let path = scratch.path("reads.fastq");
    let original = corpus(500);

    fastx::write_all(&path, &original).unwrap();
    let parsed = fastx::read_all(&path).unwrap();
    assert_eq!(parsed, original);
}

#[cfg(feature = "gzip")]
#[test]
fn gzipped_round_trip_matches_plain() {
    let scratch = Scratch::new("gz");
    let plain = scratch.path("reads.fq");
    let gzipped = scratch.path("reads.fq.gz");
    let original = corpus(300);

    fastx::write_all(&plain, &original).unwrap();
    fastx::write_all(&gzipped, &original).unwrap();

    // The gzip file must really be compressed, not just named `.gz`.
    let magic = fs::read(&gzipped).unwrap();
    assert_eq!(&magic[..2], &[0x1f, 0x8b]);

    assert_eq!(fastx::read_all(&plain).unwrap(), original);
    assert_eq!(fastx::read_all(&gzipped).unwrap(), original);
}

#[cfg(feature = "gzip")]
#[test]
fn gzip_is_detected_without_the_extension() {
    let scratch = Scratch::new("magic");
    let named = scratch.path("reads.fq.gz");
    let renamed = scratch.path("reads.fq"); // gzip bytes, misleading name
    let original = corpus(50);

    fastx::write_all(&named, &original).unwrap();
    fs::copy(&named, &renamed).unwrap();

    assert_eq!(fastx::read_all(&renamed).unwrap(), original);
}

#[cfg(feature = "zstd")]
#[test]
fn zstd_round_trip() {
    let scratch = Scratch::new("zstd");
    let plain = scratch.path("reads.fq");
    let compressed = scratch.path("reads.fq.zst");
    let original = corpus(300);

    fastx::write_all(&plain, &original).unwrap();
    fastx::write_all(&compressed, &original).unwrap();

    // Really a zstd frame, not just a file with that suffix.
    let bytes = fs::read(&compressed).unwrap();
    assert_eq!(&bytes[..4], &[0x28, 0xb5, 0x2f, 0xfd]);
    assert_eq!(fastx::read_all(&compressed).unwrap(), original);

    // Detected from the magic bytes too, so a misleading name still works.
    let misnamed = scratch.path("actually_zstd.fq");
    fs::copy(&compressed, &misnamed).unwrap();
    assert_eq!(fastx::read_all(&misnamed).unwrap(), original);

    // And it should beat gzip on size, which is the reason to reach for it.
    #[cfg(feature = "gzip")]
    {
        let gzipped = scratch.path("reads.fq.gz");
        fastx::write_all(&gzipped, &original).unwrap();
        let zstd_len = fs::metadata(&compressed).unwrap().len();
        let gzip_len = fs::metadata(&gzipped).unwrap().len();
        assert!(zstd_len < gzip_len, "zstd {zstd_len} vs gzip {gzip_len}");
    }
}

#[cfg(feature = "zstd")]
#[test]
fn zstd_cannot_be_indexed() {
    let scratch = Scratch::new("zstd-idx");
    let path = scratch.path("ref.fa.zst");
    let records: Vec<Sequence> = corpus(20).into_iter().map(Sequence::into_fasta).collect();
    fastx::write_all(&path, &records).unwrap();

    // A single frame has nothing to seek to, so indexing must refuse clearly...
    let error = FastaIndex::build_from_path(&path).unwrap_err();
    assert!(matches!(error, Error::Index(_)), "{error}");
    assert!(error.to_string().contains("Zstandard"), "{error}");
    // ...while streaming it is perfectly fine.
    assert_eq!(fastx::read_all(&path).unwrap(), records);
}

#[test]
fn fasta_round_trip_at_every_line_width() {
    let scratch = Scratch::new("widths");
    let original: Vec<Sequence> = corpus(60).into_iter().map(Sequence::into_fasta).collect();

    for width in [0usize, 1, 2, 7, 60, 80, 10_000] {
        let path = scratch.path(&format!("w{width}.fa"));
        let mut writer = WriterBuilder::new()
            .line_width(width)
            .create(&path)
            .unwrap();
        writer.write_all(&original).unwrap();
        writer.finish().unwrap();

        let parsed = fastx::read_all(&path).unwrap();
        assert_eq!(parsed, original, "line width {width}");
    }
}

#[test]
fn format_and_compression_are_inferred_from_the_path() {
    let scratch = Scratch::new("infer");
    let records = corpus(10);

    let mut names = vec![
        ("x.fa", Format::Fasta),
        ("x.fasta", Format::Fasta),
        ("x.fna", Format::Fasta),
        ("x.faa", Format::Fasta),
        ("x.fq", Format::Fastq),
        ("x.fastq", Format::Fastq),
    ];
    if cfg!(feature = "gzip") {
        names.push(("x.fna.gz", Format::Fasta));
        names.push(("x.fastq.gz", Format::Fastq));
    }

    for (name, expected) in names {
        let path = scratch.path(name);
        fastx::write_all(&path, &records).unwrap();
        let reader = fastx::open(&path).unwrap();
        assert_eq!(reader.format(), Some(expected), "{name}");
    }

    // An unknown extension has nothing to infer from.
    assert!(matches!(
        fastx::create(scratch.path("x.txt")),
        Err(Error::UnknownFormat { .. })
    ));
}

#[test]
fn writing_fastq_without_quality_is_an_error() {
    let scratch = Scratch::new("noqual");
    let path = scratch.path("out.fq");
    let records = vec![Sequence::fasta("a", b"ACGT")];
    assert!(matches!(
        fastx::write_all(&path, &records),
        Err(Error::MissingQuality { .. })
    ));
}

#[test]
fn a_single_line_chromosome_streams_fine() {
    // 4 MB on one line, read through a 4 KiB buffer: the buffer has to grow.
    let scratch = Scratch::new("chrom");
    let path = scratch.path("chrom.fa");
    let sequence: Vec<u8> = (0..4_000_000).map(|i| b"ACGT"[i % 4]).collect();
    let mut writer = WriterBuilder::new().line_width(0).create(&path).unwrap();
    writer
        .write_record(&Sequence::fasta("chrom", sequence.clone()))
        .unwrap();
    writer.finish().unwrap();

    let mut reader = fastx::ReaderBuilder::new()
        .buffer_size(4096)
        .open(&path)
        .unwrap();
    let record = reader.read_record().unwrap().expect("one record");
    assert_eq!(record.seq.len(), 4_000_000);
    assert_eq!(record.seq, sequence);
    assert!(reader.read_record().unwrap().is_none());
}

#[test]
fn malformed_input_reports_the_line() {
    // A stray line inside a FASTA record is just more sequence.
    let records = fastx::read_all_from(&b">ok\nACGT\ngarbage\n"[..]).unwrap();
    assert_eq!(records[0].seq, b"ACGTgarbage");

    // A stray line where a FASTQ header belongs is an error, on line 5.
    let err = fastx::FastxReader::new(&b"@ok\nACGT\n+\nIIII\nnot a header\n"[..])
        .nth(1)
        .unwrap()
        .unwrap_err();
    match err {
        Error::Parse { line, .. } => assert_eq!(line, 5),
        other => panic!("expected a parse error, got {other}"),
    }
    assert!(err.is_malformed());

    // A record that ends before its quality string is also an error.
    let err = fastx::FastxReader::new(&b"@ok\nACGT\n+\nIIII\n@second\nACGT\n"[..])
        .nth(1)
        .unwrap()
        .unwrap_err();
    assert!(err.is_malformed(), "{err}");
}

#[test]
fn faidx_matches_the_reader() {
    let scratch = Scratch::new("faidx");
    let path = scratch.path("ref.fa");
    let original: Vec<Sequence> = corpus(40).into_iter().map(Sequence::into_fasta).collect();

    let mut writer = WriterBuilder::new().line_width(20).create(&path).unwrap();
    writer.write_all(&original).unwrap();
    writer.finish().unwrap();

    let index = FastaIndex::build_from_path(&path).unwrap();
    index.write_to_path(&path).unwrap();
    assert_eq!(index.len(), original.len());
    assert_eq!(
        index.total_length(),
        original.iter().map(|r| r.len() as u64).sum::<u64>()
    );

    let mut fasta = IndexedFasta::open(&path).unwrap();
    for record in &original {
        let fetched = fasta.fetch(&record.id).unwrap();
        assert_eq!(fetched.seq, record.seq, "{}", record.id);

        // Every sub-region must agree with the in-memory sequence.
        for start in 0..record.len() {
            for end in [start, start + 1, record.len()] {
                if end > record.len() {
                    continue;
                }
                let region = fasta
                    .fetch_region(&record.id, start as u64, end as u64)
                    .unwrap();
                assert_eq!(
                    region.seq,
                    &record.seq[start..end],
                    "{} {start}..{end}",
                    record.id
                );
            }
        }
    }
}

#[test]
fn validation_rejects_bad_residues_at_write_time() {
    let scratch = Scratch::new("validate");
    let path = scratch.path("out.fa");
    let mut writer = WriterBuilder::new()
        .validate(Alphabet::Dna)
        .create(&path)
        .unwrap();
    assert!(writer
        .write_record(&Sequence::fasta("good", b"ACGTN"))
        .is_ok());
    assert!(matches!(
        writer.write_record(&Sequence::fasta("bad", b"ACGT?")),
        Err(Error::InvalidByte { byte: b'?', .. })
    ));
    writer.finish().unwrap();

    // The good record was still written, and it is readable.
    let parsed = fastx::read_all(&path).unwrap();
    assert_eq!(parsed.len(), 1);
    assert_eq!(parsed[0].id, "good");
}

#[test]
fn stats_of_a_file() {
    let scratch = Scratch::new("stats");
    let path = scratch.path("reads.fq");
    fastx::write_all(&path, &corpus(100)).unwrap();

    let stats = fastx::stats_of(&path).unwrap();
    assert_eq!(stats.count, 100);
    assert!(stats.mean_quality().is_some());
    assert!(stats.n50().unwrap() > 0);
    assert_eq!(
        stats.total_length,
        (0..100u64).map(|i| (i % 97) + 1).sum::<u64>()
    );
}

#[test]
fn missing_file_reports_its_name() {
    let err = fastx::open(Path::new("definitely-not-here.fa")).unwrap_err();
    assert!(err.to_string().contains("definitely-not-here.fa"), "{err}");
}

#[cfg(feature = "parallel")]
#[test]
fn parallel_and_sequential_agree() {
    use fastx::parallel;

    let scratch = Scratch::new("parallel");
    let path = scratch.path("reads.fq");
    let original = corpus(1000);
    fastx::write_all(&path, &original).unwrap();

    let sequential = fastx::stats_of(&path).unwrap();
    let parallel = parallel::par_stats_file(&path, 8).unwrap();
    assert_eq!(parallel.count, sequential.count);
    assert_eq!(parallel.total_length, sequential.total_length);
    assert_eq!(parallel.n50(), sequential.n50());
    assert_eq!(parallel.q30_fraction(), sequential.q30_fraction());

    // Chunked reads must cover every record exactly once, in order.
    let ids: Vec<String> = parallel::par_map_chunks_file(&path, 8, |reader| {
        let mut ids = Vec::new();
        reader.for_each_record(|r| {
            ids.push(r.id.clone());
            Ok(())
        })?;
        Ok(ids)
    })
    .unwrap()
    .into_iter()
    .flatten()
    .collect();
    let expected: Vec<String> = original.iter().map(|r| r.id.clone()).collect();
    assert_eq!(ids, expected);
}