qubit-json 0.12.2

// =============================================================================
//    Copyright (c) 2026 Haixing Hu.
//
//    SPDX-License-Identifier: Apache-2.0
//
//    Licensed under the Apache License, Version 2.0.
// =============================================================================
//! Tests reusable JSON encode sessions.

use std::io;
use std::io::Write;
use std::panic::AssertUnwindSafe;
use std::panic::catch_unwind;

use qubit_budget::ResourceBudget;
use qubit_budget::ResourceLimit;
use qubit_budget::StructureLimits;
use qubit_budget::json::JsonEncodeLimits;
use qubit_budget::json::JsonEncodeSession;
use qubit_budget::json::JsonMeasurement;
use qubit_budget::json::JsonResource;
use qubit_budget::json::JsonValueBudget;
use qubit_budget::json::JsonValueLimits;
use serde::Serialize;
use serde::Serializer;
use serde::ser::Error as SerError;
use serde::ser::SerializeSeq;

use crate::encode::json_encode_test_support::encode;
use crate::encode::json_encode_test_support::write_buffered;
use crate::encode::json_encode_test_support::write_incremental;

/// Value that serializes one accepted prefix before returning an error.
struct FailsAfterPrefix;

impl Serialize for FailsAfterPrefix {
    /// Emits one sequence item, then returns a custom serialization error.
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut sequence = serializer.serialize_seq(Some(2))?;
        sequence.serialize_element(&1_u8)?;
        Err(SerError::custom("intentional streaming failure"))
    }
}

/// Value that serializes one accepted prefix before unwinding.
struct PanicsAfterPrefix;

impl Serialize for PanicsAfterPrefix {
    /// Emits one sequence item, then panics during serialization.
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
    where
        S: Serializer,
    {
        let mut sequence = serializer.serialize_seq(Some(2))?;
        sequence.serialize_element(&1_u8)?;
        panic!("intentional streaming panic");
    }
}

/// Writer that accepts a bounded prefix, then fails every later write.
struct PrefixWriter {
    /// Bytes accepted before the configured failure boundary.
    accepted: Vec<u8>,
    /// Maximum number of bytes that may be accepted.
    maximum: usize,
}

impl Write for PrefixWriter {
    /// Accepts only the remaining prefix capacity before returning an I/O
    /// error.
    fn write(&mut self, input: &[u8]) -> io::Result<usize> {
        if self.accepted.len() == self.maximum {
            return Err(io::Error::other("intentional writer failure"));
        }
        let accepted = (self.maximum - self.accepted.len()).min(input.len());
        self.accepted.extend_from_slice(&input[..accepted]);
        Ok(accepted)
    }

    /// Does not add buffering beyond the accepted prefix.
    fn flush(&mut self) -> io::Result<()> {
        Ok(())
    }
}

/// Verifies encode sessions expose their output-byte resource.
#[test]
fn test_encode_session_exposes_only_output_resource() {
    let encode = JsonEncodeSession::from_limits(
        JsonEncodeLimits::<JsonResource, usize>::builder()
            .output_bytes_limit(ResourceLimit::new(JsonResource::OutputBytes, 8))
            .build(),
    );

    assert_eq!(encode.max_output_bytes(), Some(8));
}

/// Verifies output-byte consumption is cumulative and atomic within one
/// attempt.
#[test]
fn test_encode_attempt_consumes_output_bytes_atomically() {
    let mut session = JsonEncodeSession::from_limits(
        JsonEncodeLimits::<JsonResource, usize>::builder()
            .output_bytes_limit(ResourceLimit::new(JsonResource::OutputBytes, 3))
            .build(),
    );

    let mut attempt = session.begin_value();
    attempt
        .try_consume_output_bytes(3)
        .expect("exact output fits");
    let error = attempt
        .try_consume_output_bytes(1)
        .expect_err("output budget is exhausted");
    assert_eq!(error.resource(), &JsonResource::OutputBytes);
}

/// Verifies encode attempts preserve every embedded JSON value limit.
#[test]
fn test_encode_attempt_preserves_embedded_value_limits() {
    let value_limits = JsonValueLimits::<JsonResource, usize>::builder()
        .string_bytes_limit(ResourceLimit::new(JsonResource::StringBytes, 2))
        .payload_bytes_limit(ResourceLimit::new(JsonResource::PayloadBytes, 3))
        .structure_limits(
            StructureLimits::builder().nodes_limit(ResourceLimit::new(JsonResource::Nodes, 2)),
        )
        .build();
    let mut session = JsonEncodeSession::from_limits(
        JsonEncodeLimits::<JsonResource, usize>::builder()
            .value_limits(value_limits)
            .build(),
    );

    let mut attempt = session.begin_value();
    attempt
        .try_admit(JsonMeasurement::String { depth: 1, bytes: 2 })
        .expect("exact string limit fits");
    let string_error = attempt
        .try_admit(JsonMeasurement::String { depth: 1, bytes: 3 })
        .expect_err("overlong string poisons the attempt");
    assert_eq!(string_error.resource(), &JsonResource::StringBytes);
    let commit_error = attempt
        .commit()
        .expect_err("poisoned string attempt cannot commit");
    assert_eq!(commit_error.resource(), string_error.resource());

    let mut attempt = session.begin_value();
    attempt
        .try_admit(JsonMeasurement::String { depth: 1, bytes: 2 })
        .expect("string payload fits");
    attempt
        .try_admit(JsonMeasurement::Key { bytes: 1 })
        .expect("exact payload limit fits");
    let payload_error = attempt
        .try_admit(JsonMeasurement::Key { bytes: 1 })
        .expect_err("exhausted payload poisons the attempt");
    assert_eq!(payload_error.resource(), &JsonResource::PayloadBytes);
    let commit_error = attempt
        .commit()
        .expect_err("poisoned payload attempt cannot commit");
    assert_eq!(commit_error.resource(), payload_error.resource());

    let mut attempt = session.begin_value();
    attempt
        .try_admit(JsonMeasurement::Null { depth: 1 })
        .expect("first node fits");
    attempt
        .try_admit(JsonMeasurement::Null { depth: 1 })
        .expect("exact node limit fits");
    let node_error = attempt
        .try_admit(JsonMeasurement::Null { depth: 1 })
        .expect_err("exhausted node budget poisons the attempt");
    assert_eq!(node_error.resource(), &JsonResource::Nodes);
    let commit_error = attempt
        .commit()
        .expect_err("poisoned node attempt cannot commit");
    assert_eq!(commit_error.resource(), node_error.resource());
    assert_eq!(session.value_budget().used_nodes(), Some(0));
    assert_eq!(session.value_budget().used_payload_bytes(), Some(0));
}

/// Verifies borrowed encode sessions charge caller-owned budgets in place.
#[test]
fn test_encode_session_borrows_output_and_value_budgets() {
    let mut output = ResourceBudget::new(JsonResource::OutputBytes, 32);
    let mut value = JsonValueBudget::new(
        JsonValueLimits::<JsonResource, usize>::builder()
            .structure_limits(
                StructureLimits::builder().nodes_limit(ResourceLimit::new(JsonResource::Nodes, 16)),
            )
            .build(),
    );
    let mut session = JsonEncodeSession::borrowing_output(&mut output, &mut value);

    let encoded = encode(&serde_json::json!({"name": "qubit"}), &mut session)
        .expect("borrowed budgets should support online encoding");

    assert_eq!(encoded, br#"{"name":"qubit"}"#);
    assert_eq!(output.used(), encoded.len());
    assert!(
        value
            .used_nodes()
            .expect("nodes limit should be configured")
            > 0
    );
}

/// Verifies buffered vector serialization publishes neither output nor value
/// accounting when a streamed Serde value fails.
#[test]
fn test_encode_stream_failure_rolls_back_output_and_value() {
    let mut session = JsonEncodeSession::from_limits(
        JsonEncodeLimits::<JsonResource, usize>::builder()
            .max_output_bytes(64)
            .max_nodes(2)
            .build(),
    );

    assert!(encode(&FailsAfterPrefix, &mut session).is_err());
    assert_eq!(session.output_budget().expect("output budget").used(), 0,);
    assert_eq!(session.value_budget().used_nodes(), Some(0));
}

/// Verifies buffered writes retain their externally accepted prefix and output
/// charge, while their staged value accounting rolls back after I/O failure.
#[test]
fn test_buffered_writer_failure_keeps_prefix_and_rolls_back_value() {
    let mut session = JsonEncodeSession::from_limits(
        JsonEncodeLimits::<JsonResource, usize>::builder()
            .max_output_bytes(64)
            .max_nodes(3)
            .build(),
    );
    let mut writer = PrefixWriter {
        accepted: Vec::new(),
        maximum: 2,
    };

    assert!(write_buffered(&mut writer, &[1_u8, 2_u8], &mut session).is_err());
    assert_eq!(writer.accepted.len(), 2);
    assert_eq!(
        session.output_budget().expect("output budget").used(),
        writer.accepted.len(),
    );
    assert_eq!(session.value_budget().used_nodes(), Some(0));
}

/// Verifies incremental streaming retains its accepted output prefix but drops
/// staged values when serialization returns an error.
#[test]
fn test_incremental_stream_failure_keeps_prefix_and_rolls_back_value() {
    let mut session = JsonEncodeSession::from_limits(
        JsonEncodeLimits::<JsonResource, usize>::builder()
            .max_output_bytes(64)
            .max_nodes(2)
            .build(),
    );
    let mut writer = Vec::new();

    assert!(write_incremental(&mut writer, &FailsAfterPrefix, &mut session).is_err());
    assert!(!writer.is_empty());
    assert_eq!(
        session.output_budget().expect("output budget").used(),
        writer.len(),
    );
    assert_eq!(session.value_budget().used_nodes(), Some(0));
}

/// Verifies panic unwind keeps incremental output side effects but rolls back
/// staged values so the session can encode a later value.
#[test]
fn test_incremental_panic_keeps_prefix_and_reuses_value_capacity() {
    let mut session = JsonEncodeSession::from_limits(
        JsonEncodeLimits::<JsonResource, usize>::builder()
            .max_output_bytes(64)
            .max_nodes(2)
            .build(),
    );
    let mut writer = Vec::new();

    let result = catch_unwind(AssertUnwindSafe(|| {
        write_incremental(&mut writer, &PanicsAfterPrefix, &mut session)
    }));

    assert!(result.is_err());
    assert!(!writer.is_empty());
    assert_eq!(
        session.output_budget().expect("output budget").used(),
        writer.len(),
    );
    assert_eq!(session.value_budget().used_nodes(), Some(0));
    encode(&true, &mut session).expect("panic must leave value capacity for the next encode");
    assert_eq!(session.value_budget().used_nodes(), Some(1));
}

mod incremental_writer_failures {
    use std::io;
    use std::io::Write;

    use qubit_budget::json::JsonEncodeLimits;
    use qubit_budget::json::JsonResource;
    use qubit_io::ByteOutput;
    use qubit_json::encode::JsonEncodeErrorKind;
    use qubit_json::encode::JsonEncoder;

    #[derive(Clone, Copy)]
    enum WriterMode {
        FailAfterThree,
        InterruptedOnce,
        WriteZero,
        FlushError,
    }

    struct ControlledWriter {
        bytes: Vec<u8>,
        accepted: usize,
        mode: WriterMode,
        interrupted: bool,
    }

    impl ControlledWriter {
        #[must_use]
        fn new(mode: WriterMode) -> Self {
            Self {
                bytes: Vec::new(),
                accepted: 0,
                mode,
                interrupted: false,
            }
        }
    }

    impl Write for ControlledWriter {
        fn write(&mut self, input: &[u8]) -> io::Result<usize> {
            match self.mode {
                WriterMode::FailAfterThree if self.accepted >= 3 => {
                    return Err(io::Error::other("controlled failure"));
                }
                WriterMode::InterruptedOnce if !self.interrupted => {
                    self.interrupted = true;
                    return Err(io::Error::from(io::ErrorKind::Interrupted));
                }
                WriterMode::WriteZero => return Ok(0),
                _ => {}
            }

            let remaining = match self.mode {
                WriterMode::FailAfterThree => 3 - self.accepted,
                _ => usize::MAX,
            };
            let count = input.len().min(2).min(remaining);
            self.bytes.extend_from_slice(&input[..count]);
            self.accepted += count;
            Ok(count)
        }

        fn flush(&mut self) -> io::Result<()> {
            if matches!(self.mode, WriterMode::FlushError) {
                Err(io::Error::other("controlled flush failure"))
            } else {
                Ok(())
            }
        }
    }

    fn encoder_with_output_limit(limit: usize) -> JsonEncoder<'static> {
        JsonEncoder::with_limits(
            JsonEncodeLimits::<JsonResource, usize>::builder()
                .max_output_bytes(limit)
                .build(),
        )
    }

    #[test]
    fn test_incremental_writer_charges_only_accepted_prefix() {
        let mut encoder = encoder_with_output_limit(64);
        let mut writer = ControlledWriter::new(WriterMode::FailAfterThree);
        let output: &mut dyn ByteOutput = &mut writer;
        let error = encoder.write_incremental(output, &"abcdef").unwrap_err();

        assert_eq!(error.kind(), JsonEncodeErrorKind::Write);
        assert_eq!(writer.bytes, b"\"ab");
        assert_eq!(encoder.session().output_budget().unwrap().used(), 3);
    }

    #[test]
    fn test_incremental_writer_retries_interrupted() {
        let mut encoder = JsonEncoder::unlimited();
        let mut writer = ControlledWriter::new(WriterMode::InterruptedOnce);
        let output: &mut dyn ByteOutput = &mut writer;
        encoder.write_incremental(output, &true).unwrap();

        assert_eq!(writer.bytes, b"true");
    }

    #[test]
    fn test_incremental_writer_reports_write_zero() {
        let mut encoder = JsonEncoder::unlimited();
        let mut writer = ControlledWriter::new(WriterMode::WriteZero);
        let output: &mut dyn ByteOutput = &mut writer;
        let error = encoder.write_incremental(output, &true).unwrap_err();

        assert_eq!(error.kind(), JsonEncodeErrorKind::Write);
        assert_eq!(
            error.write_error().unwrap().kind(),
            io::ErrorKind::WriteZero
        );
        assert!(writer.bytes.is_empty());
    }

    #[test]
    fn test_incremental_writer_reports_flush_failure_and_charges_output() {
        let mut encoder = JsonEncoder::unlimited();
        let mut writer = ControlledWriter::new(WriterMode::FlushError);
        let output: &mut dyn ByteOutput = &mut writer;
        let error = encoder.write_incremental(output, &true).unwrap_err();

        assert_eq!(error.kind(), JsonEncodeErrorKind::Write);
        assert_eq!(error.write_error().unwrap().kind(), io::ErrorKind::Other);
        assert_eq!(writer.bytes, b"true");
    }
}