vix.ing · top · new · best · stats

Automated Assertion Generation and Regression Testing for Machine Learning Notebooks

2025/09/17 by Yingao Elaine Yao, Yingao, Yao +7
Computer Science · Decision Sciences · #Adversarial Robustness in Machine Learning #FOS: Computer and information sciences #Scientific Computing and Data Management #Software Engineering (cs.SE) #Software Testing and Debugging Techniques #cs.SE

paper · pdf · doi:10.48550/arxiv.2509.13656

22 pages, 2 figures, 6 tables

openalex publication_date 2025/09/17 · openalex created_date 2025/10/10 · openalex updated_date 2026/07/28 · arxiv created 2026/08/05 · arxiv updated 2026/08/06

Abstract

Jupyter Notebooks have become the de-facto choice for data scientists and machine learning (ML) engineers for prototyping and experimenting with ML pipelines, offering a rich interactive interface for code, data, and visualization in one place. However, notebooks provide limited support for testing. As a result, during continuous development, many silent (non-crashing) regressions often go unnoticed, making notebooks unreliable and results hard to reproduce. We introduce NBTestGen, the first automated assertion generation approach for ML notebooks. NBTestGen generates regression-based assertions that check properties of data processing, model building, and model evaluation steps in a typical ML notebook. To support integration of such assertions in notebooks, we introduce the first regression testing framework (called NBTest) that can be used as a Jupyter plugin and allows developers to write cell-scoped assertions in notebooks. Such assertions are 1) cell-scoped: they are linked to specific notebook cells and execute only after those cells are executed, 2) non-intrusive: they do not block notebook execution (in a Jupyter session), so that development can continue when they fail, and 3) they integrate with pytest and CI pipelines, allowing developers to easily do regression testing of their notebooks. We evaluate NBTestGen on a corpus of 585 notebooks from the popular Kaggle platform. NBTestGen generates a total of 21542 assertions (36.82 on average per notebook). The generated assertions kill 72.21% of ML-specific mutations, while maintaining a high passrate of 100.00%. We also show that NBTestGen can detect 69.68% of historical regressions in 531 older versions of Kaggle notebooks. A popular ML library, SHAP, integrated NBTest into their CI. Further, we perform a user study with 17 ML developers that shows that such users find NBTest highly intuitive and useful.

Citations

Related