owlen/crates/tools/notebook/tests/notebook_tests.rs

use tools_notebook::*;
use std::fs;
use tempfile::tempdir;

#[test]
fn notebook_round_trip_preserves_metadata() {
    let dir = tempdir().unwrap();
    let notebook_path = dir.path().join("test.ipynb");

    // Create a sample notebook with metadata
    let notebook_json = r##"{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "source": ["print('hello world')"],
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": ["# Test Notebook", "This is a test."]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.0"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}"##;

    fs::write(&notebook_path, notebook_json).unwrap();

    // Read the notebook
    let notebook = read_notebook(&notebook_path).unwrap();

    // Verify structure
    assert_eq!(notebook.cells.len(), 2);
    assert_eq!(notebook.nbformat, 4);
    assert_eq!(notebook.nbformat_minor, 5);

    // Verify metadata
    assert!(notebook.metadata.kernelspec.is_some());
    let kernelspec = notebook.metadata.kernelspec.as_ref().unwrap();
    assert_eq!(kernelspec.language, "python");
    assert_eq!(kernelspec.name, "python3");

    assert!(notebook.metadata.language_info.is_some());
    let lang_info = notebook.metadata.language_info.as_ref().unwrap();
    assert_eq!(lang_info.name, "python");
    assert_eq!(lang_info.version, Some("3.9.0".to_string()));

    // Write it back
    let output_path = dir.path().join("output.ipynb");
    write_notebook(&output_path, &notebook).unwrap();

    // Read it again
    let notebook2 = read_notebook(&output_path).unwrap();

    // Verify metadata is preserved
    assert_eq!(notebook2.nbformat, 4);
    assert_eq!(notebook2.nbformat_minor, 5);
    assert!(notebook2.metadata.kernelspec.is_some());
    assert_eq!(
        notebook2.metadata.kernelspec.as_ref().unwrap().language,
        "python"
    );
}

#[test]
fn notebook_edit_cell_content() {
    let dir = tempdir().unwrap();
    let notebook_path = dir.path().join("test.ipynb");

    let notebook_json = r##"{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "source": ["x = 1"],
      "outputs": []
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "source": ["y = 2"],
      "outputs": []
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
}"##;

    fs::write(&notebook_path, notebook_json).unwrap();

    let mut notebook = read_notebook(&notebook_path).unwrap();

    // Edit the first cell
    edit_notebook(
        &mut notebook,
        NotebookEdit::EditCell {
            index: 0,
            source: vec!["x = 10\n".to_string(), "print(x)\n".to_string()],
        },
    )
    .unwrap();

    // Verify the edit
    assert_eq!(notebook.cells[0].source.len(), 2);
    assert_eq!(notebook.cells[0].source[0], "x = 10\n");
    assert_eq!(notebook.cells[0].source[1], "print(x)\n");

    // Second cell should be unchanged
    assert_eq!(notebook.cells[1].source[0], "y = 2");
}

#[test]
fn notebook_add_delete_cells() {
    let dir = tempdir().unwrap();
    let notebook_path = dir.path().join("test.ipynb");

    let notebook_json = r##"{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "source": ["x = 1"],
      "outputs": []
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
}"##;

    fs::write(&notebook_path, notebook_json).unwrap();

    let mut notebook = read_notebook(&notebook_path).unwrap();
    assert_eq!(notebook.cells.len(), 1);

    // Add a cell at the end
    let new_cell = new_code_cell(vec!["y = 2\n".to_string()]);
    edit_notebook(
        &mut notebook,
        NotebookEdit::AddCell {
            index: 1,
            cell: new_cell,
        },
    )
    .unwrap();

    assert_eq!(notebook.cells.len(), 2);
    assert_eq!(notebook.cells[1].source[0], "y = 2\n");

    // Add a cell at the beginning
    let first_cell = new_markdown_cell(vec!["# Header\n".to_string()]);
    edit_notebook(
        &mut notebook,
        NotebookEdit::AddCell {
            index: 0,
            cell: first_cell,
        },
    )
    .unwrap();

    assert_eq!(notebook.cells.len(), 3);
    assert_eq!(notebook.cells[0].cell_type, "markdown");
    assert_eq!(notebook.cells[0].source[0], "# Header\n");
    assert_eq!(notebook.cells[1].source[0], "x = 1"); // Original first cell is now second

    // Delete the middle cell
    edit_notebook(&mut notebook, NotebookEdit::DeleteCell { index: 1 }).unwrap();

    assert_eq!(notebook.cells.len(), 2);
    assert_eq!(notebook.cells[0].cell_type, "markdown");
    assert_eq!(notebook.cells[1].source[0], "y = 2\n");
}

#[test]
fn notebook_edit_out_of_bounds() {
    let dir = tempdir().unwrap();
    let notebook_path = dir.path().join("test.ipynb");

    let notebook_json = r##"{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "source": ["x = 1\n"],
      "outputs": []
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
}"##;

    fs::write(&notebook_path, notebook_json).unwrap();

    let mut notebook = read_notebook(&notebook_path).unwrap();

    // Try to edit non-existent cell
    let result = edit_notebook(
        &mut notebook,
        NotebookEdit::EditCell {
            index: 5,
            source: vec!["bad\n".to_string()],
        },
    );

    assert!(result.is_err());
    assert!(result.unwrap_err().to_string().contains("out of bounds"));
}

#[test]
fn notebook_with_outputs_preserved() {
    let dir = tempdir().unwrap();
    let notebook_path = dir.path().join("test.ipynb");

    let notebook_json = r##"{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "source": ["print('hello')\n"],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": ["hello\n"]
        }
      ]
    }
  ],
  "metadata": {},
  "nbformat": 4,
  "nbformat_minor": 5
}"##;

    fs::write(&notebook_path, notebook_json).unwrap();

    let notebook = read_notebook(&notebook_path).unwrap();

    assert_eq!(notebook.cells[0].outputs.len(), 1);
    assert_eq!(notebook.cells[0].outputs[0].output_type, "stream");

    // Write and read back
    let output_path = dir.path().join("output.ipynb");
    write_notebook(&output_path, &notebook).unwrap();

    let notebook2 = read_notebook(&output_path).unwrap();
    assert_eq!(notebook2.cells[0].outputs.len(), 1);
    assert_eq!(notebook2.cells[0].outputs[0].output_type, "stream");
}

#[test]
fn cell_source_as_string_concatenates() {
    let cell = new_code_cell(vec![
        "import numpy as np\n".to_string(),
        "arr = np.array([1, 2, 3])\n".to_string(),
    ]);

    let source = cell_source_as_string(&cell);
    assert_eq!(source, "import numpy as np\narr = np.array([1, 2, 3])\n");
}