| 1 | package builtin |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "os" |
| 7 | "path/filepath" |
| 8 | "strings" |
| 9 | "testing" |
| 10 | ) |
| 11 | |
| 12 | // sampleNotebook is a minimal but realistic .ipynb: a markdown cell and a code |
| 13 | // cell with an output + execution_count, plus top-level metadata/nbformat that |
| 14 | // must survive a round-trip. |
| 15 | const sampleNotebook = `{ |
| 16 | "cells": [ |
| 17 | {"cell_type": "markdown", "id": "intro", "metadata": {}, "source": ["# Title\n", "text"]}, |
| 18 | {"cell_type": "code", "id": "c1", "metadata": {}, "execution_count": 5, "outputs": [{"output_type": "stream", "text": "old"}], "source": ["print(1)\n"]} |
| 19 | ], |
| 20 | "metadata": {"kernelspec": {"name": "python3"}}, |
| 21 | "nbformat": 4, |
| 22 | "nbformat_minor": 5 |
| 23 | }` |
| 24 | |
| 25 | func writeNotebook(t *testing.T) string { |
| 26 | t.Helper() |
| 27 | dir := t.TempDir() |
| 28 | p := filepath.Join(dir, "nb.ipynb") |
| 29 | if err := os.WriteFile(p, []byte(sampleNotebook), 0o644); err != nil { |
| 30 | t.Fatal(err) |
| 31 | } |
| 32 | return p |
| 33 | } |
| 34 | |
| 35 | func runNotebookEdit(t *testing.T, path string, args map[string]any) (string, error) { |
| 36 | t.Helper() |
| 37 | args["path"] = path |
| 38 | raw, _ := json.Marshal(args) |
| 39 | return notebookEdit{}.Execute(context.Background(), raw) |
| 40 | } |
| 41 | |
| 42 | func readCells(t *testing.T, path string) []map[string]json.RawMessage { |
| 43 | t.Helper() |
| 44 | data, err := os.ReadFile(path) |
| 45 | if err != nil { |
| 46 | t.Fatal(err) |
| 47 | } |
| 48 | nb, err := parseNotebook(data) |
| 49 | if err != nil { |
| 50 | t.Fatalf("result is not valid notebook JSON: %v", err) |
| 51 | } |
| 52 | return nb.cells |
| 53 | } |
| 54 | |
| 55 | func TestNotebookReplaceBySource(t *testing.T) { |
| 56 | p := writeNotebook(t) |
| 57 | if _, err := runNotebookEdit(t, p, map[string]any{"cell_number": 1, "new_source": "print(42)\n"}); err != nil { |
| 58 | t.Fatal(err) |
| 59 | } |
| 60 | cells := readCells(t, p) |
| 61 | if len(cells) != 2 { |
| 62 | t.Fatalf("replace changed cell count: %d", len(cells)) |
| 63 | } |
| 64 | if got := string(cells[1]["source"]); !strings.Contains(got, "print(42)") { |
| 65 | t.Errorf("source not replaced: %s", got) |
| 66 | } |
| 67 | // Editing a code cell clears its outputs + execution_count. |
| 68 | if got := string(cells[1]["outputs"]); got != "[]" { |
| 69 | t.Errorf("outputs not cleared: %s", got) |
| 70 | } |
| 71 | if got := string(cells[1]["execution_count"]); got != "null" { |
| 72 | t.Errorf("execution_count not cleared: %s", got) |
| 73 | } |
| 74 | } |
| 75 | |
| 76 | func TestNotebookRetypeNormalizesOutputs(t *testing.T) { |
| 77 | p := writeNotebook(t) |
| 78 | if _, err := runNotebookEdit(t, p, map[string]any{"cell_number": 1, "cell_type": "markdown", "new_source": "# now md"}); err != nil { |
| 79 | t.Fatal(err) |
| 80 | } |
| 81 | md := readCells(t, p)[1] |
| 82 | if _, has := md["outputs"]; has { |
| 83 | t.Errorf("code→markdown left 'outputs' (invalid nbformat): %s", md["outputs"]) |
| 84 | } |
| 85 | if _, has := md["execution_count"]; has { |
| 86 | t.Errorf("code→markdown left 'execution_count' (invalid nbformat): %s", md["execution_count"]) |
| 87 | } |
| 88 | |
| 89 | if _, err := runNotebookEdit(t, p, map[string]any{"cell_number": 0, "cell_type": "code", "new_source": "y = 2\n"}); err != nil { |
| 90 | t.Fatal(err) |
| 91 | } |
| 92 | code := readCells(t, p)[0] |
| 93 | if string(code["outputs"]) != "[]" || string(code["execution_count"]) != "null" { |
| 94 | t.Errorf("markdown→code missing output scaffolding: outputs=%s exec=%s", code["outputs"], code["execution_count"]) |
| 95 | } |
| 96 | } |
| 97 | |
| 98 | func TestNotebookReplaceByID(t *testing.T) { |
| 99 | p := writeNotebook(t) |
| 100 | if _, err := runNotebookEdit(t, p, map[string]any{"cell_id": "intro", "new_source": "# New"}); err != nil { |
| 101 | t.Fatal(err) |
| 102 | } |
| 103 | cells := readCells(t, p) |
| 104 | if got := string(cells[0]["source"]); !strings.Contains(got, "# New") { |
| 105 | t.Errorf("cell_id target not replaced: %s", got) |
| 106 | } |
| 107 | } |
| 108 | |
| 109 | func TestNotebookInsertAfter(t *testing.T) { |
| 110 | p := writeNotebook(t) |
| 111 | if _, err := runNotebookEdit(t, p, map[string]any{"edit_mode": "insert", "cell_number": 0, "cell_type": "code", "new_source": "x = 1\n"}); err != nil { |
| 112 | t.Fatal(err) |
| 113 | } |
| 114 | cells := readCells(t, p) |
| 115 | if len(cells) != 3 { |
| 116 | t.Fatalf("insert should add a cell, got %d", len(cells)) |
| 117 | } |
| 118 | if got := string(cells[1]["cell_type"]); got != `"code"` { |
| 119 | t.Errorf("inserted cell type wrong: %s", got) |
| 120 | } |
| 121 | if got := string(cells[1]["source"]); !strings.Contains(got, "x = 1") { |
| 122 | t.Errorf("inserted source wrong: %s", got) |
| 123 | } |
| 124 | // A code cell gets outputs/execution_count scaffolding. |
| 125 | if _, ok := cells[1]["outputs"]; !ok { |
| 126 | t.Error("inserted code cell missing outputs") |
| 127 | } |
| 128 | } |
| 129 | |
| 130 | func TestNotebookInsertPrepend(t *testing.T) { |
| 131 | p := writeNotebook(t) |
| 132 | if _, err := runNotebookEdit(t, p, map[string]any{"edit_mode": "insert", "cell_number": -1, "cell_type": "markdown", "new_source": "top"}); err != nil { |
| 133 | t.Fatal(err) |
| 134 | } |
| 135 | cells := readCells(t, p) |
| 136 | if got := string(cells[0]["source"]); !strings.Contains(got, "top") { |
| 137 | t.Errorf("prepend should land at index 0: %s", got) |
| 138 | } |
| 139 | } |
| 140 | |
| 141 | func TestNotebookDelete(t *testing.T) { |
| 142 | p := writeNotebook(t) |
| 143 | if _, err := runNotebookEdit(t, p, map[string]any{"edit_mode": "delete", "cell_number": 0}); err != nil { |
| 144 | t.Fatal(err) |
| 145 | } |
| 146 | cells := readCells(t, p) |
| 147 | if len(cells) != 1 { |
| 148 | t.Fatalf("delete should leave 1 cell, got %d", len(cells)) |
| 149 | } |
| 150 | if got := cellID(cells[0]); got != "c1" { |
| 151 | t.Errorf("wrong cell deleted; remaining id = %q", got) |
| 152 | } |
| 153 | } |
| 154 | |
| 155 | func TestNotebookPreservesTopLevelKeys(t *testing.T) { |
| 156 | p := writeNotebook(t) |
| 157 | if _, err := runNotebookEdit(t, p, map[string]any{"cell_number": 0, "new_source": "# x"}); err != nil { |
| 158 | t.Fatal(err) |
| 159 | } |
| 160 | data, _ := os.ReadFile(p) |
| 161 | var top map[string]json.RawMessage |
| 162 | if err := json.Unmarshal(data, &top); err != nil { |
| 163 | t.Fatal(err) |
| 164 | } |
| 165 | for _, k := range []string{"metadata", "nbformat", "nbformat_minor"} { |
| 166 | if _, ok := top[k]; !ok { |
| 167 | t.Errorf("top-level key %q lost on round-trip", k) |
| 168 | } |
| 169 | } |
| 170 | } |
| 171 | |
| 172 | func TestNotebookErrors(t *testing.T) { |
| 173 | p := writeNotebook(t) |
| 174 | if _, err := runNotebookEdit(t, p, map[string]any{"cell_number": 9, "new_source": "x"}); err == nil { |
| 175 | t.Error("out-of-range cell_number should error") |
| 176 | } |
| 177 | if _, err := runNotebookEdit(t, p, map[string]any{"cell_id": "nope", "edit_mode": "delete"}); err == nil { |
| 178 | t.Error("unknown cell_id should error") |
| 179 | } |
| 180 | if _, err := runNotebookEdit(t, p, map[string]any{"edit_mode": "insert", "cell_number": 0, "new_source": "x"}); err == nil { |
| 181 | t.Error("insert without cell_type should error") |
| 182 | } |
| 183 | if _, err := runNotebookEdit(t, p, map[string]any{"edit_mode": "bogus"}); err == nil { |
| 184 | t.Error("bad edit_mode should error") |
| 185 | } |
| 186 | } |
| 187 | |
| 188 | // TestNotebookPreviewMatchesExecute checks the Previewer mirrors Execute: the |
| 189 | // previewed NewText equals the file content Execute persists. |
| 190 | func TestNotebookPreviewMatchesExecute(t *testing.T) { |
| 191 | p := writeNotebook(t) |
| 192 | args := map[string]any{"path": p, "cell_number": 1, "new_source": "print(99)\n"} |
| 193 | raw, _ := json.Marshal(args) |
| 194 | |
| 195 | change, err := notebookEdit{}.Preview(raw) |
| 196 | if err != nil { |
| 197 | t.Fatalf("preview: %v", err) |
| 198 | } |
| 199 | if _, err := (notebookEdit{}).Execute(context.Background(), raw); err != nil { |
| 200 | t.Fatalf("execute: %v", err) |
| 201 | } |
| 202 | persisted, _ := os.ReadFile(p) |
| 203 | if change.NewText != string(persisted) { |
| 204 | t.Errorf("preview NewText != persisted content:\npreview:\n%s\npersisted:\n%s", change.NewText, persisted) |
| 205 | } |
| 206 | } |
| 207 | |
| 208 | // TestNotebookContentAlias accepts the write_file-style "content" field as an |
| 209 | // alias for new_source, and defaults to the only cell when no target is given — |
| 210 | // the near-miss shape a model reaches for, which should succeed not loop. |
| 211 | func TestNotebookContentAlias(t *testing.T) { |
| 212 | dir := t.TempDir() |
| 213 | p := filepath.Join(dir, "one.ipynb") |
| 214 | one := `{"cells":[{"cell_type":"code","id":"a","metadata":{},"execution_count":null,"outputs":[],"source":["print('hi')\n"]}],"metadata":{},"nbformat":4,"nbformat_minor":5}` |
| 215 | if err := os.WriteFile(p, []byte(one), 0o644); err != nil { |
| 216 | t.Fatal(err) |
| 217 | } |
| 218 | // {content, path} with no cell_number — the exact shape that looped before. |
| 219 | if _, err := runNotebookEdit(t, p, map[string]any{"content": "print(\"world\")\n"}); err != nil { |
| 220 | t.Fatalf("content-alias single-cell replace should succeed, got: %v", err) |
| 221 | } |
| 222 | cells := readCells(t, p) |
| 223 | if got := string(cells[0]["source"]); !strings.Contains(got, `print(\"world\")`) { |
| 224 | t.Errorf("alias source not applied: %s", got) |
| 225 | } |
| 226 | } |
| 227 | |
| 228 | // TestNotebookMissingTargetMultiCell still requires an explicit target when the |
| 229 | // notebook is ambiguous (more than one cell), with an instructive message. |
| 230 | func TestNotebookMissingTargetMultiCell(t *testing.T) { |
| 231 | p := writeNotebook(t) // 2 cells |
| 232 | _, err := runNotebookEdit(t, p, map[string]any{"new_source": "x"}) |
| 233 | if err == nil { |
| 234 | t.Fatal("multi-cell replace with no target should error") |
| 235 | } |
| 236 | if !strings.Contains(err.Error(), "cell_number") { |
| 237 | t.Errorf("error should name cell_number: %v", err) |
| 238 | } |
| 239 | } |
| 240 | |
| 241 | // TestNotebookSourceLines checks the nbformat line-array encoding: \n kept on |
| 242 | // every line but the last. |
| 243 | func TestNotebookSourceLines(t *testing.T) { |
| 244 | if got := string(sourceLines("a\nb\n")); got != `["a\n","b\n"]` { |
| 245 | t.Errorf("source line encoding = %s", got) |
| 246 | } |
| 247 | if got := string(sourceLines("solo")); got != `["solo"]` { |
| 248 | t.Errorf("single line = %s", got) |
| 249 | } |
| 250 | if got := string(sourceLines("")); got != `[]` { |
| 251 | t.Errorf("empty = %s", got) |
| 252 | } |
| 253 | } |
| 254 |