{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6c55f6fb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:20.087561Z",
     "iopub.status.busy": "2026-06-19T02:32:20.087380Z",
     "iopub.status.idle": "2026-06-19T02:32:25.503208Z",
     "shell.execute_reply": "2026-06-19T02:32:25.502356Z"
    },
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
     ]
    }
   ],
   "source": [
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import brainmass\n",
    "import brainstate\n",
    "import brainunit as u\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import numpy as np\n",
    "brainstate.environ.set(dt=0.1 * u.ms)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bc7acce",
   "metadata": {},
   "source": [
    "# Testing\n",
    "\n",
    "Testing guidelines for `brainmass`.\n",
    "\n",
    "## Running Tests\n",
    "\n",
    "Run all tests:\n",
    "\n",
    "```bash\n",
    "pytest tests/\n",
    "```\n",
    "\n",
    "Run specific test file:\n",
    "\n",
    "```bash\n",
    "pytest tests/test_models.py\n",
    "```\n",
    "\n",
    "Run with coverage:\n",
    "\n",
    "```bash\n",
    "pytest --cov=brainmass tests/\n",
    "```\n",
    "\n",
    "## Test Structure\n",
    "\n",
    "```text\n",
    "tests/\n",
    "├── test_models.py          # Neural mass model tests\n",
    "├── test_noise.py           # Noise process tests\n",
    "├── test_coupling.py        # Coupling mechanism tests\n",
    "├── test_forward.py         # Forward model tests\n",
    "└── test_utils.py           # Utility function tests\n",
    "```\n",
    "\n",
    "## Writing Tests\n",
    "\n",
    "### Basic Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d0edb0af",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:25.507284Z",
     "iopub.status.busy": "2026-06-19T02:32:25.506552Z",
     "iopub.status.idle": "2026-06-19T02:32:25.513233Z",
     "shell.execute_reply": "2026-06-19T02:32:25.512042Z"
    }
   },
   "outputs": [],
   "source": [
    "import pytest\n",
    "import jax.numpy as jnp\n",
    "import brainmass\n",
    "\n",
    "def test_hopf_oscillator():\n",
    "    # Create model\n",
    "    model = brainmass.HopfStep(in_size=10, w=0.2)\n",
    "\n",
    "    # Initialize\n",
    "    model.init_all_states()\n",
    "\n",
    "    # Test update\n",
    "    output = model.update()\n",
    "\n",
    "    # Assertions\n",
    "    assert output.shape == (10,)\n",
    "    assert not jnp.any(jnp.isnan(output))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69b732a6",
   "metadata": {},
   "source": [
    "### Parametrized Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6bd5a264",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:25.516665Z",
     "iopub.status.busy": "2026-06-19T02:32:25.516287Z",
     "iopub.status.idle": "2026-06-19T02:32:25.522768Z",
     "shell.execute_reply": "2026-06-19T02:32:25.520982Z"
    }
   },
   "outputs": [],
   "source": [
    "@pytest.mark.parametrize(\"in_size\", [1, 10, 100])\n",
    "def test_model_shapes(in_size):\n",
    "    model = brainmass.WilsonCowanStep(in_size=in_size)\n",
    "    model.init_all_states()\n",
    "\n",
    "    output = model.update(rE_inp=0.5, rI_inp=0.2)\n",
    "    assert output.shape == (in_size,)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ef5bf7f",
   "metadata": {},
   "source": [
    "### Batch Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "641b6cc4",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-06-19T02:32:25.527201Z",
     "iopub.status.busy": "2026-06-19T02:32:25.526687Z",
     "iopub.status.idle": "2026-06-19T02:32:25.533065Z",
     "shell.execute_reply": "2026-06-19T02:32:25.531999Z"
    }
   },
   "outputs": [],
   "source": [
    "@pytest.mark.parametrize(\"batch_size\", [None, 1, 32])\n",
    "def test_batching(batch_size):\n",
    "    model = brainmass.HopfStep(in_size=5, w=0.2)\n",
    "    model.init_all_states(batch_size=batch_size)\n",
    "\n",
    "    output = model.update()\n",
    "\n",
    "    if batch_size is None:\n",
    "        assert output.shape == (5,)\n",
    "    else:\n",
    "        assert output.shape == (batch_size, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5259e77b",
   "metadata": {},
   "source": [
    "## Test Best Practices\n",
    "\n",
    "1. **Test one thing per test**\n",
    "2. **Use descriptive test names**\n",
    "3. **Check shapes, values, and edge cases**\n",
    "4. **Test with and without batch dimension**\n",
    "5. **Test with units where applicable**\n",
    "\n",
    "## Coverage\n",
    "\n",
    "Aim for >80% code coverage:\n",
    "\n",
    "```bash\n",
    "pytest --cov=brainmass --cov-report=html tests/\n",
    "# Open htmlcov/index.html\n",
    "```\n",
    "\n",
    "## Continuous Integration\n",
    "\n",
    "Tests run automatically on:\n",
    "\n",
    "- Every pull request\n",
    "- Pushes to main branch\n",
    "- Nightly builds\n",
    "\n",
    "Ensure tests pass before submitting PR.\n",
    "\n",
    "## See Also\n",
    "\n",
    "- pytest documentation\n",
    "- {doc}`contributing` - Contribution process"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
