From d9eee7be9bc1a79d130f5204a52ec8bc28c4724e Mon Sep 17 00:00:00 2001
From: Heberto Mayorquin <h.mayorquin@gmail.com>
Date: Sat, 30 Mar 2024 14:01:38 -0600
Subject: [PATCH] changelog, ruff and test

---
 CHANGELOG.md                      |  1 +
 src/pynwb/testing/mock/ecephys.py | 11 ++++-------
 tests/unit/test_mock.py           |  4 ++++
 3 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8a7515c3e..b3a79ca2d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@
   - Support `stimulus_template` as optional predefined column in `IntracellularStimuliTable`. @stephprince [#1815](https://github.com/NeurodataWithoutBorders/pynwb/pull/1815)
   - Support `NWBDataInterface` and `DynamicTable` in `NWBFile.stimulus`. @rly [#1842](https://github.com/NeurodataWithoutBorders/pynwb/pull/1842)
 - Added support for python 3.12 and upgraded dependency versions. This also includes infrastructure updates for developers. @mavaylon1 [#1853](https://github.com/NeurodataWithoutBorders/pynwb/pull/1853)
+- Added `mock_Units` for generating Units tables:  @h-mayorquin [#1853](https://github.com/NeurodataWithoutBorders/pynwb/pull/1853)
 
 ### Bug fixes
 - Fix bug with reading file with linked `TimeSeriesReferenceVectorData` @rly [#1865](https://github.com/NeurodataWithoutBorders/pynwb/pull/1865)
diff --git a/src/pynwb/testing/mock/ecephys.py b/src/pynwb/testing/mock/ecephys.py
index bc64e49ed..6ea86c31c 100644
--- a/src/pynwb/testing/mock/ecephys.py
+++ b/src/pynwb/testing/mock/ecephys.py
@@ -9,6 +9,7 @@
 from ...ecephys import ElectricalSeries, ElectrodeGroup, SpikeEventSeries
 from .device import mock_Device
 from .utils import name_generator
+from ...misc import Units
 
 
 def mock_ElectrodeGroup(
@@ -121,11 +122,7 @@ def mock_SpikeEventSeries(
     return spike_event_series
 
 
-from pynwb.misc import Units
-import numpy as np 
-
-
-def mock_Units(num_units: int = 10, max_spikes: int = 10, seed: Optional[int] = None):
+def mock_Units(num_units: int = 10, max_spikes_per_unit: int = 10, seed: Optional[int] = None):
 
     units_table = Units()
     units_table.add_column(
@@ -136,8 +133,8 @@ def mock_Units(num_units: int = 10, max_spikes: int = 10, seed: Optional[int] =
         seed = 0
     rng = np.random.default_rng(seed=seed)
 
-    times = rng.random(size=(num_units, max_spikes)).cumsum(axis=1)
-    spikes_per_unit = rng.integers(1, max_spikes, size=num_units)
+    times = rng.random(size=(num_units, max_spikes_per_unit)).cumsum(axis=1)
+    spikes_per_unit = rng.integers(1, max_spikes_per_unit, size=num_units)
 
     spike_times = []
     for unit_index in range(num_units):
diff --git a/tests/unit/test_mock.py b/tests/unit/test_mock.py
index 72174f018..af2fc7dc9 100644
--- a/tests/unit/test_mock.py
+++ b/tests/unit/test_mock.py
@@ -35,6 +35,7 @@
     mock_ElectrodeTable,
     mock_ElectricalSeries,
     mock_SpikeEventSeries,
+    mock_Units,
 )
 
 from pynwb.testing.mock.icephys import (
@@ -82,6 +83,7 @@
     mock_IntracellularElectrode,
     mock_CurrentClampStimulusSeries,
     mock_IntracellularRecordingsTable,
+    mock_Units,
 ]
 
 
@@ -119,3 +121,5 @@ def test_name_generator():
 
     assert name_generator("TimeSeries") == "TimeSeries"
     assert name_generator("TimeSeries") == "TimeSeries2"
+
+