From 111d0dc504161a7735bb6941be66fd0aeacdadee Mon Sep 17 00:00:00 2001
From: LorenzzoQM <lorenzzo.quevedomantovani@colorado.edu>
Date: Tue, 21 Nov 2023 15:04:37 -0700
Subject: [PATCH 1/4] setDataBuffer method added to dataStorageUnitBase

---
 .../dataStorageUnitBase.cpp                   | 58 +++++++++++++++++--
 .../_GeneralModuleFiles/dataStorageUnitBase.h |  9 +--
 2 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp
index 550e72d8e6..059cc487d2 100644
--- a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp
+++ b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp
@@ -27,8 +27,12 @@
 DataStorageUnitBase::DataStorageUnitBase(){
     this->previousTime = 0; //! - previousTime initialized to 0.
     this->nodeDataUseInMsgs.clear(); //! - Clear the vector of input messages.
-    this->storedDataSum = 0.0; //! - Initialize the dataSum to 0.
+    this->storedDataSum = 0; //! - Initialize the dataSum to 0.
     this->netBaud = 0.0; //! - Initialize the netBaudRate to 0.
+    //! - Zero out the partitions
+    for(uint64_t i = 0; i < this->storedData.size(); i++){
+        this->storedData[i].dataInstanceSum = 0;
+    }
     return;
 }
 
@@ -168,18 +172,18 @@ void DataStorageUnitBase::integrateDataStatus(double currentTime){
         index = messageInStoredData(&(*it));
 
         //! - If the storage capacity has not been reached or the baudRate is less than 0 and won't take below 0, then add the data
-       if ((this->storedDataSum < this->storageCapacity) || (it->baudRate < 0)) {
+       if ((this->storedDataSum + round(it->baudRate * this->currentTimestep) <= this->storageCapacity) || (it->baudRate < 0)) {
            //! - if a dataNode exists in storedData vector, integrate and add to current amount
            if (index != -1) {
                //! Only perform if this operation will not take the sum below zero
                if ((this->storedData[(size_t) index].dataInstanceSum + it->baudRate * (this->currentTimestep)) >= 0) {
-                   this->storedData[(size_t) index].dataInstanceSum += it->baudRate * (this->currentTimestep);
+                   this->storedData[(size_t) index].dataInstanceSum += round(it->baudRate * (this->currentTimestep));
                }
                //! - if a dataNode does not exist in storedData, add it to storedData, integrate baud rate, and add amount
            }
            else if (strcmp(it->dataName, "") != 0) {
                strncpy(tmpDataInstance.dataInstanceName, it->dataName, sizeof(tmpDataInstance.dataInstanceName));
-               tmpDataInstance.dataInstanceSum = it->baudRate * (this->currentTimestep);
+               tmpDataInstance.dataInstanceSum = round(it->baudRate * (this->currentTimestep));
                this->storedData.push_back(tmpDataInstance);
            }
        }
@@ -214,8 +218,8 @@ int DataStorageUnitBase::messageInStoredData(DataNodeUsageMsgPayload *tmpNodeMsg
 /*! Sums all of the data in the storedData vector
  @return double
  */
-double DataStorageUnitBase::sumAllData(){
-    double dataSum = 0.0;
+long long int DataStorageUnitBase::sumAllData(){
+    double dataSum = 0;
 
     std::vector<dataInstance>::iterator it;
     for(it = storedData.begin(); it != storedData.end(); it++) {
@@ -248,3 +252,45 @@ bool DataStorageUnitBase::customReadMessages()
 {
     return true;
 }
+
+/*! Adds a specific amount of data to the storedData vector once
+ @param partitionName //Name of the partition to add data to
+ @param data          //Amount of data to add to the partition
+ @return void
+ */
+void DataStorageUnitBase::setDataBuffer(std::string partitionName, long long int data)
+{
+    dataInstance tmpDataInstance;
+
+    int index = -1;
+    for (uint64_t i = 0; i < this->storedData.size(); i++){
+        if (strcmp(this->storedData[i].dataInstanceName, partitionName.c_str()) == 0){
+            index = (int) i;
+        }
+    }
+
+    //! - If the new data won't overflow the storage capacity, then add the data
+    if (this->storedDataSum + data <= this->storageCapacity) {
+        //! - if a dataNode exists in storedData vector, integrate and add to current amount
+        if (index != -1) {
+            //! Only perform if this operation will not take the sum below zero
+            if ((this->storedData[(size_t) index].dataInstanceSum + data) >= 0) {
+                this->storedData[(size_t) index].dataInstanceSum += data;
+            }
+            
+        }
+        //! - if a dataNode does not exist in storedData, add it to storedData, and add amount
+        else if (strcmp(partitionName.c_str(), "") != 0) {
+            strncpy(tmpDataInstance.dataInstanceName, partitionName.c_str(), sizeof(tmpDataInstance.dataInstanceName));
+            //! Only perform this operation if the resulting sum in the partition is not negative. If it is, initialize to zero.
+            if (data < 0) {
+                data = 0;
+            }
+            tmpDataInstance.dataInstanceSum = data;
+            this->storedData.push_back(tmpDataInstance);
+        }
+    }
+
+    //! - Sum all data in storedData vector
+    this->storedDataSum = this->sumAllData();
+}
\ No newline at end of file
diff --git a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.h b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.h
index 3b5a073981..a466b93dea 100644
--- a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.h
+++ b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.h
@@ -34,7 +34,7 @@
 
 struct dataInstance {
     char dataInstanceName[128];     //!< data instance name
-    double dataInstanceSum;         //!< data instance sum value, bits
+    long long int dataInstanceSum;         //!< data instance sum value, bits
 }; //!< Struct for instances of data stored in a buffer. Includes names and amounts.
 
 /*! @brief on-board data handling base class */
@@ -45,6 +45,7 @@ class DataStorageUnitBase: public SysModel {
     void Reset(uint64_t CurrentSimNanos);
     void addDataNodeToModel(Message<DataNodeUsageMsgPayload> *tmpNodeMsg); //!< Adds dataNode to the storageUnit
     void UpdateState(uint64_t CurrentSimNanos);
+    void setDataBuffer(std::string partitionName, long long int data); //!< Adds/removes the data from the partitionName partition once
 
 protected:
     void writeMessages(uint64_t CurrentClock);
@@ -54,18 +55,18 @@ class DataStorageUnitBase: public SysModel {
     virtual void customWriteMessages(uint64_t CurrentClock); //!< custom Write method, similar to customSelfInit.
     virtual bool customReadMessages(); //!< Custom read method, similar to customSelfInit; returns `true' by default.
     int messageInStoredData(DataNodeUsageMsgPayload *tmpNodeMsg); //!< Returns index of the dataName if it's already in storedData
-    double sumAllData(); //!< Sums all of the data in the storedData vector
+    long long int sumAllData(); //!< Sums all of the data in the storedData vector
 
 public:
     std::vector<ReadFunctor<DataNodeUsageMsgPayload>> nodeDataUseInMsgs; //!< Vector of data node input message names
     Message<DataStorageStatusMsgPayload> storageUnitDataOutMsg; //!< Vector of message names to be written out by the storage unit
-    double storageCapacity; //!< Storage capacity of the storage unit
+    long long int storageCapacity; //!< Storage capacity of the storage unit
     BSKLogger bskLogger;    //!< logging variable
 
 protected:
     DataStorageStatusMsgPayload storageStatusMsg; //!< class variable
     std::vector<DataNodeUsageMsgPayload> nodeBaudMsgs; //!< class variable
-    double storedDataSum; //!< [bits] Stored data in bits.
+    long long int storedDataSum; //!< [bits] Stored data in bits.
     std::vector<dataInstance> storedData; //!< Vector of data. Represents the makeup of the data buffer.
     double previousTime; //!< Previous time used for integration
     double currentTimestep;//!< [s] Timestep duration in seconds.

From 757847927e678c49ff74a05ae4536fedeb66f301 Mon Sep 17 00:00:00 2001
From: LorenzzoQM <lorenzzo.quevedomantovani@colorado.edu>
Date: Fri, 13 Oct 2023 12:12:03 -0600
Subject: [PATCH 2/4] Method added to simple and partitioned StorageUnit

---
 .../dataStorageUnitBase.cpp                   |  5 -----
 .../storageUnit/partitionedStorageUnit.cpp    | 21 +++++++++++++++----
 .../storageUnit/partitionedStorageUnit.h      |  1 +
 .../storageUnit/partitionedStorageUnit.i      | 18 ++++++++++++++++
 .../storageUnit/simpleStorageUnit.cpp         | 19 ++++++++++++-----
 .../storageUnit/simpleStorageUnit.h           |  1 +
 .../storageUnit/simpleStorageUnit.i           |  7 +++++++
 7 files changed, 58 insertions(+), 14 deletions(-)

diff --git a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp
index 059cc487d2..8b5464cb45 100644
--- a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp
+++ b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.cpp
@@ -53,11 +53,6 @@ void DataStorageUnitBase::Reset(uint64_t CurrentSimNanos)
 {
     this->previousTime = 0;
 
-    //! - Zero out the partitions
-    for(uint64_t i = 0; i < this->storedData.size(); i++){
-        this->storedData[i].dataInstanceSum = 0.0;
-    }
-
     //! - call the custom environment module reset method
     customReset(CurrentSimNanos);
 
diff --git a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.cpp b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.cpp
index 1db18bac3f..86c3da297f 100644
--- a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.cpp
+++ b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.cpp
@@ -28,8 +28,8 @@
  @return void;
  */
 PartitionedStorageUnit::PartitionedStorageUnit(){
-    this->storageCapacity = -1.0;
-    this->storedDataSum = 0.0;
+    this->storageCapacity = 0;
+    this->storedDataSum = 0;
     return;
 }
 
@@ -45,7 +45,7 @@ PartitionedStorageUnit::~PartitionedStorageUnit(){
  @return void
  */
 void PartitionedStorageUnit::customReset(uint64_t currentClock){
-    if (this->storageCapacity <= 0.0) {
+    if (this->storageCapacity <= 0) {
         bskLogger.bskLog(BSK_INFORMATION, "The storageCapacity variable must be set to a positive value.");
     }
     return;
@@ -58,7 +58,20 @@ void PartitionedStorageUnit::customReset(uint64_t currentClock){
 void PartitionedStorageUnit::addPartition(std::string dataName){
     dataInstance tmpDataInstance;
     strncpy(tmpDataInstance.dataInstanceName, dataName.c_str(), sizeof(tmpDataInstance.dataInstanceName));
-    tmpDataInstance.dataInstanceSum = 0.0;
+    tmpDataInstance.dataInstanceSum = 0;
     this->storedData.push_back(tmpDataInstance);
     return;
+}
+
+/*! Adds a specific amount of data to the specified partitions once
+ @param partitionNames  //Vector of partition names
+ @param data            //Vector of data to be added to each partition in partitionNames
+ @return void
+ */
+void PartitionedStorageUnit::setDataBuffer(std::vector<std::string> partitionNames, std::vector<long long int> data){
+
+    for (int i = 0; i < partitionNames.size(); i++)
+    {
+        PartitionedStorageUnit::DataStorageUnitBase::setDataBuffer(partitionNames[i], data[i]);
+    }
 }
\ No newline at end of file
diff --git a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.h b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.h
index 7a453b0ce6..d7814563a5 100644
--- a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.h
+++ b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.h
@@ -30,6 +30,7 @@ class PartitionedStorageUnit: public DataStorageUnitBase {
     PartitionedStorageUnit();
     ~PartitionedStorageUnit();
     void addPartition(std::string dataName);
+    void setDataBuffer(std::vector<std::string> partitionNames, std::vector<long long int> data); //!< Adds/removes the data from the partitionNames partitions
 
 private:
     void customReset(uint64_t CurrentClock) override;
diff --git a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.i b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.i
index 674738f97a..2c5810fdf9 100644
--- a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.i
+++ b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.i
@@ -32,7 +32,25 @@ from Basilisk.architecture.swig_common_model import *
 %include "sys_model.i"
 %include "std_vector.i"
 %include "cstring.i"
+%include "swig_common_model.i"
+%include "stdint.i"
 
+//When using scientific notation in Python (1E9), it is interpreted as float
+// giving a type error when assigning storageCapacity or adding data through
+// setDataBuffer. This maps that float and vector of floats to long int in 
+// C++ in this module.
+%typemap(in) long long int {
+    $1 = static_cast<long long int>(PyFloat_AsDouble($input));
+}
+%typemap(in) std::vector<long long int> (std::vector<long long int> temp) {
+    size_t size = PyList_Size($input);
+    temp.reserve(size);
+    for (size_t i = 0; i < size; ++i) {
+        PyObject* item = PyList_GetItem($input, i);
+        temp.push_back(static_cast<long long int>(PyFloat_AsDouble(item)));
+    }
+    $1 = temp;
+}
 
 %include "simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.h"
 %include "partitionedStorageUnit.h"
diff --git a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.cpp b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.cpp
index 246128e1c1..761f975b13 100644
--- a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.cpp
+++ b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.cpp
@@ -27,8 +27,8 @@
  @return void
  */
 SimpleStorageUnit::SimpleStorageUnit(){
-    this->storageCapacity = -1.0;
-    this->storedDataSum = 0.0;
+    this->storageCapacity = 0;
+    this->storedDataSum = 0;
     return;
 }
 
@@ -43,7 +43,7 @@ SimpleStorageUnit::~SimpleStorageUnit(){
  @param currentClock
  */
 void SimpleStorageUnit::customReset(uint64_t currentClock){
-    if (this->storageCapacity <= 0.0) {
+    if (this->storageCapacity <= 0) {
         bskLogger.bskLog(BSK_INFORMATION, "The storageCapacity variable must be set to a positive value.");
     }
     return;
@@ -63,10 +63,10 @@ void SimpleStorageUnit::integrateDataStatus(double currentTime){
         if (storedData.size() == 0){
             this->storedData.push_back({{'S','T','O','R','E','D',' ','D','A','T','A'}, 0});
         }
-        else if ((this->storedDataSum < this->storageCapacity) || (it->baudRate <= 0)){
+        else if ((this->storedDataSum + round(it->baudRate * this->currentTimestep) < this->storageCapacity) || (it->baudRate <= 0)){
             //! - Only perform the operation if it will not result in less than 0 data
             if ((this->storedData[0].dataInstanceSum + it->baudRate * (this->currentTimestep)) >= 0){
-                this->storedData[0].dataInstanceSum += it->baudRate * (this->currentTimestep);
+                this->storedData[0].dataInstanceSum += round(it->baudRate * (this->currentTimestep));
             }
         }
         this->netBaud += it->baudRate;
@@ -79,3 +79,12 @@ void SimpleStorageUnit::integrateDataStatus(double currentTime){
     this->previousTime = currentTime;
     return;
 }
+
+/*! Adds a specific amount of data to the storedData vector once
+ @param data //Data to be added to the "STORED DATA" partition
+ @return void
+ */
+void SimpleStorageUnit::setDataBuffer(long long int data){
+    std::string partitionName = "STORED DATA";
+    SimpleStorageUnit::DataStorageUnitBase::setDataBuffer(partitionName, data);
+}
\ No newline at end of file
diff --git a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.h b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.h
index 55ade06695..e64deedd1b 100644
--- a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.h
+++ b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.h
@@ -30,6 +30,7 @@ class SimpleStorageUnit: public DataStorageUnitBase {
 public:
     SimpleStorageUnit();
     ~SimpleStorageUnit();
+    void setDataBuffer(long long int data); //!< Method to add/remove data from the storage unit once
 
 private:
     void customReset(uint64_t CurrentClock); //!< Custom Reset method
diff --git a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.i b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.i
index 65b1870cd3..f5c0345bba 100644
--- a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.i
+++ b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.i
@@ -33,6 +33,13 @@ from Basilisk.architecture.swig_common_model import *
 %include "sys_model.i"
 %include "stdint.i"
 
+//When using scientific notation in Python (1E9), it is interpreted as float
+// giving a type error when assigning storageCapacity or using setDataBuffer.
+// This maps that float to long int in C++ in this module.
+%typemap(in) long long int {
+    $1 = static_cast<long long int>(PyFloat_AsDouble($input));
+}
+
 %include "simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.h"
 %include "simpleStorageUnit.h"
 %include "architecture/msgPayloadDefC/DataNodeUsageMsgPayload.h"

From 672b908c80f56c1d8ab39217a8d0594fe4562bda Mon Sep 17 00:00:00 2001
From: LorenzzoQM <lorenzzo.quevedomantovani@colorado.edu>
Date: Tue, 21 Nov 2023 15:06:44 -0700
Subject: [PATCH 3/4] Updated tests for setDataBuffer method

---
 .../_UnitTest/test_partitionedStorageUnit.py  | 263 +++++++++++++++---
 .../_UnitTest/test_simpleStorageUnit.py       | 186 ++++++++++---
 2 files changed, 376 insertions(+), 73 deletions(-)

diff --git a/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_partitionedStorageUnit.py b/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_partitionedStorageUnit.py
index a39a90d36e..955f8c440b 100644
--- a/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_partitionedStorageUnit.py
+++ b/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_partitionedStorageUnit.py
@@ -18,6 +18,8 @@
 
 import inspect
 import os
+import pytest
+import numpy as np
 
 filename = inspect.getframeinfo(inspect.currentframe()).filename
 path = os.path.dirname(os.path.abspath(filename))
@@ -26,31 +28,29 @@
 
 # Import all of the modules that we are going to be called in this simulation
 from Basilisk.utilities import SimulationBaseClass
-from Basilisk.utilities import unitTestSupport                  # general support file with common unit test functions
 from Basilisk.simulation import partitionedStorageUnit
 from Basilisk.architecture import messaging
 from Basilisk.utilities import macros
 
-def test_module(show_plots):
-    # each test method requires a single assert method to be called
-    [testResults, testMessage] = check_storage_limits(show_plots)
-    assert testResults < 1, testMessage
 
+params_storage_limits = [(1200, 1200, 2400, 2400),
+                     (600, 1200, 3600, 3600),
+                     (600, 600, 10000, 6000)]
 
-def check_storage_limits(show_plots):
+@pytest.mark.parametrize("baudRate_1, baudRate_2, storageCapacity, expectedStorage",
+                          params_storage_limits)
+def test_storage_limits(baudRate_1, baudRate_2, storageCapacity, expectedStorage):
     """
     Tests:
 
-    1. Whether the partitionedStorageUnit can add multiple nodes (core base class functionality);
-    2. That the partitionedStorageUnit correctly evaluates how much stored data it should have given a pair of
-       1200 baud input messages.
+    1. Whether the partitionedStorageUnit can add multiple nodes (core base class 
+        functionality);
+    2. That the partitionedStorageUnit correctly evaluates how much stored data it 
+        should have given a pair of baud input messages.
 
-    :param show_plots: Not used; no plots to be shown.
     :return:
     """
 
-    testFailCount = 0                       # zero unit test result counter
-    testMessages = []                       # create empty array to store test log messages
     unitTaskName = "unitTask"               # arbitrary name (don't change)
     unitProcessName = "TestProcess"         # arbitrary name (don't change)
 
@@ -63,15 +63,15 @@ def check_storage_limits(show_plots):
     testProc.addTask(unitTestSim.CreateNewTask(unitTaskName, testProcessRate))
 
     test_storage_unit = partitionedStorageUnit.PartitionedStorageUnit()
-    test_storage_unit.storageCapacity = 2400. # bit capacity.
+    test_storage_unit.storageCapacity = storageCapacity # bit capacity.
 
     dataMsg1 = messaging.DataNodeUsageMsgPayload()
-    dataMsg1.baudRate = 1200. # baud
+    dataMsg1.baudRate = baudRate_1 # baud
     dataMsg1.dataName = "node_1_msg"
     dat1Msg = messaging.DataNodeUsageMsg().write(dataMsg1)
 
     dataMsg2 = messaging.DataNodeUsageMsgPayload()
-    dataMsg2.baudRate = 1200. # baud
+    dataMsg2.baudRate = baudRate_2 # baud
     dataMsg2.dataName = "node_2_msg"
     dat2Msg = messaging.DataNodeUsageMsg().write(dataMsg2)
 
@@ -93,36 +93,227 @@ def check_storage_limits(show_plots):
     capacityLog = dataLog.storageCapacity
     netBaudLog = dataLog.currentNetBaud
 
-    #   Check 1 - is net baud rate equal to 2400.?
+    #   Check 1 - is net baud rate correct?
     for ind in range(0,len(netBaudLog)):
         currentBaud = netBaudLog[ind]
-        if currentBaud !=2400.:
-            testFailCount +=1
-            testMessages.append("FAILED: PartitionedStorageUnit did not correctly log the net baud rate.")
+        np.testing.assert_allclose(currentBaud, baudRate_1 + baudRate_2, atol=1e-1,
+            err_msg=("FAILED: PartitionedStorageUnit did not correctly log baud rate."))
 
-    #print(netBaudLog)
+    #   Check 2 - is used storage space correct?
+    np.testing.assert_allclose(storedDataLog[-1], expectedStorage, atol=1e-4,
+        err_msg=("FAILED: PartitionedStorageUnit did not track integrated data."))
 
-    if not unitTestSupport.isDoubleEqualRelative((2400.),storedDataLog[-1], 1e-8):
-        testFailCount+=1
-        testMessages.append("FAILED: PartitionedStorageUnit did not track integrated data. Returned "+str(storedDataLog[-1,1])+", expected "+str((2400.)))
+    #   Check 3 - is the amount of data more than zero and less than the capacity?
+    for ind in range(0,len(storedDataLog)):
+        assert storedDataLog[ind] <= capacityLog[ind] or np.isclose(storedDataLog[ind], 
+            capacityLog[ind]), (
+                "FAILED: PartitionedStorageUnit's stored data exceeded its capacity.")
+
+        assert storedDataLog[ind] >= 0., (
+            "FAILED: PartitionedStorageUnit's stored data was negative.")
+
+
+params_set_data = [(1200, 1200, ['test'], [1200], 2400, 2400),
+                   (600, 600, ['test'], [0], 10000, 6000),
+                   (600, 600, ['test'], [4000], 10000, 10000),
+                   (0, 0, ['test'], [1000], 2000, 1000),
+                   (0, 0, ['test'], [-1000], 2000, 0),
+                   (1000, 0, ['node_1_msg'], [-2000], 6000, 3000),
+                   (0, 0, ['test'], [3000], 2000, 0),
+                   (600, 0, ['node_1_msg'], [3000], 10000, 6000),
+                   (300, 600, ['node_1_msg'], [3000], 10000, 7500),
+                   (600, 600, ['node_1_msg', 'node_2_msg'], [1000, 1000], 10000, 8000),
+                   (600, 300, ['test', 'node_2_msg'], [1000, 1000], 10000, 6500)]
+
+@pytest.mark.parametrize(
+    "baudRate_1, baudRate_2, partitionName, add_data, storageCapacity, expectedStorage",
+    params_set_data)
+def test_set_data_buffer(baudRate_1, baudRate_2, partitionName, 
+                         add_data, storageCapacity, expectedStorage):
+    """
+    Tests:
+
+    1. Whether the partitionedStorageUnit properly adds data in different partitions 
+        using the setDataBuffer method;
+    2. That the partitionedStorageUnit correctly evaluates how much stored data it 
+        should have given a pair of input messages and using setDataBuffer.
+
+    :return:
+    """
+
+    unitTaskName = "unitTask"               # arbitrary name (don't change)
+    unitProcessName = "TestProcess"         # arbitrary name (don't change)
+
+    # Create a sim module as an empty container
+    unitTestSim = SimulationBaseClass.SimBaseClass()
+
+    # Create test thread
+    testProcessRate = macros.sec2nano(0.1)     # update process rate update time
+    testProc = unitTestSim.CreateNewProcess(unitProcessName)
+    testProc.addTask(unitTestSim.CreateNewTask(unitTaskName, testProcessRate))
+
+    test_storage_unit = partitionedStorageUnit.PartitionedStorageUnit()
+    test_storage_unit.storageCapacity = storageCapacity # bit capacity.
+
+    dataMsg1 = messaging.DataNodeUsageMsgPayload()
+    dataMsg1.baudRate = baudRate_1 # baud
+    dataMsg1.dataName = "node_1_msg"
+    dat1Msg = messaging.DataNodeUsageMsg().write(dataMsg1)
+
+    dataMsg2 = messaging.DataNodeUsageMsgPayload()
+    dataMsg2.baudRate = baudRate_2 # baud
+    dataMsg2.dataName = "node_2_msg"
+    dat2Msg = messaging.DataNodeUsageMsg().write(dataMsg2)
+
+    # Test the addNodeToStorage method:
+    test_storage_unit.addDataNodeToModel(dat1Msg)
+    test_storage_unit.addDataNodeToModel(dat2Msg)
+
+    unitTestSim.AddModelToTask(unitTaskName, test_storage_unit)
+
+    dataLog = test_storage_unit.storageUnitDataOutMsg.recorder()
+    unitTestSim.AddModelToTask(unitTaskName, dataLog)
 
-    #print(storedDataLog)
+    # Initialize the partition
+    initData = [0 for i in range(0, len(partitionName))]
+    test_storage_unit.setDataBuffer(partitionName, initData)
+
+    unitTestSim.InitializeSimulation()
+    sim_time = 5.0
+    unitTestSim.ConfigureStopTime(macros.sec2nano(sim_time - 1.0))
+    unitTestSim.ExecuteSimulation()
+
+    test_storage_unit.setDataBuffer(partitionName, add_data)
+    unitTestSim.ConfigureStopTime(macros.sec2nano(sim_time))
+    unitTestSim.ExecuteSimulation()
 
+    storedDataLog = dataLog.storageLevel
+    capacityLog = dataLog.storageCapacity
+    netBaudLog = dataLog.currentNetBaud
+
+    #   Check 1 - is net baud rate correct?
+    for ind in range(0,len(netBaudLog)):
+        currentBaud = netBaudLog[ind]
+        np.testing.assert_allclose(currentBaud, baudRate_1 + baudRate_2, atol=1e-4,
+            err_msg=("FAILED: PartitionedStorageUnit did not correctly log baud rate."))
+
+    #   Check 2 - is used storage space correct?
+    np.testing.assert_allclose(storedDataLog[-1], expectedStorage, atol=1e-4,
+        err_msg=("FAILED: PartitionedStorageUnit did not track integrated data."))
+
+    #   Check 3 - is the amount of data more than zero and less than the capacity?
     for ind in range(0,len(storedDataLog)):
-        if storedDataLog[ind] > capacityLog[ind]:
-            testFailCount +=1
-            testMessages.append("FAILED: PartitionedStorageUnit's stored data exceeded its capacity.")
+        assert storedDataLog[ind] <= capacityLog[ind] or np.isclose(storedDataLog[ind], 
+            capacityLog[ind]), (
+                "FAILED: PartitionedStorageUnit's stored data exceeded its capacity.")
+
+        assert storedDataLog[ind] >= 0., (
+            "FAILED: PartitionedStorageUnit's stored data was negative.")
+
+
+params_set_data = [(600, 600, ['test'], [0], 1E4),
+    (0, 0, ['test'], [1000], 2000),
+    (1000, 0, ['node_1_msg'], [-2000], 6000),
+    (600, 0, ['node_1_msg'], [3000], 1e4),
+    (600, 600, ['node_1_msg'], [3000], 10000),
+    (600, 600, ['node_1_msg', 'node_2_msg'], [1e3, 1e3], 10000),
+    (600, 300, ['test', 'node_1_msg', 'node_2_msg'], [1500, 1e3, 1000], 10000)]
+
+@pytest.mark.parametrize(
+    "baudRate_1, baudRate_2, partitionName, add_data, storageCapacity",
+    params_set_data)
+def test_set_data_buffer_partition(baudRate_1, baudRate_2, partitionName, 
+                                   add_data, storageCapacity):
+
+    """
+    Tests:
+
+    1. Whether the partitionedStorageUnit manages the data in already existing 
+        and new partitions correctly;
+
+    :return:
+    """
+
+    unitTaskName = "unitTask"               # arbitrary name (don't change)
+    unitProcessName = "TestProcess"         # arbitrary name (don't change)
+
+    # Create a sim module as an empty container
+    unitTestSim = SimulationBaseClass.SimBaseClass()
+
+    # Create test thread
+    testProcessRate = macros.sec2nano(0.1)     # update process rate update time
+    testProc = unitTestSim.CreateNewProcess(unitProcessName)
+    testProc.addTask(unitTestSim.CreateNewTask(unitTaskName, testProcessRate))
+
+    test_storage_unit = partitionedStorageUnit.PartitionedStorageUnit()
+    test_storage_unit.storageCapacity = storageCapacity # bit capacity.
+
+    dataMsg1 = messaging.DataNodeUsageMsgPayload()
+    dataMsg1.baudRate = baudRate_1 # baud
+    dataMsg1.dataName = "node_1_msg"
+    dat1Msg = messaging.DataNodeUsageMsg().write(dataMsg1)
+
+    dataMsg2 = messaging.DataNodeUsageMsgPayload()
+    dataMsg2.baudRate = baudRate_2 # baud
+    dataMsg2.dataName = "node_2_msg"
+    dat2Msg = messaging.DataNodeUsageMsg().write(dataMsg2)
+
+    # Test the addNodeToStorage method:
+    test_storage_unit.addDataNodeToModel(dat1Msg)
+    test_storage_unit.addDataNodeToModel(dat2Msg)
+
+    unitTestSim.AddModelToTask(unitTaskName, test_storage_unit)
+
+    dataLog = test_storage_unit.storageUnitDataOutMsg.recorder()
+    unitTestSim.AddModelToTask(unitTaskName, dataLog)
+
+    # Initialize the partition
+    initData = [0 for i in range(0, len(partitionName))]
+    test_storage_unit.setDataBuffer(partitionName, initData)
+
+    unitTestSim.InitializeSimulation()
+    sim_time = 5.0
+    unitTestSim.ConfigureStopTime(macros.sec2nano(sim_time - 1.0))
+    unitTestSim.ExecuteSimulation()
+
+    test_storage_unit.setDataBuffer(partitionName, add_data)
+    unitTestSim.ConfigureStopTime(macros.sec2nano(sim_time))
+    unitTestSim.ExecuteSimulation()
+
+    dataNameVec = dataLog.storedDataName
+    dataVec = dataLog.storedData
+
+    # Check 1 - are the partition names in the data name vector?
+    for partName in partitionName:
+        assert partName in list(dataNameVec[-1]), (
+            "FAILED: PartitionedStorageUnit did not add the new partition.")
 
-        if storedDataLog[ind] < 0.:
-            testFailCount +=1
-            testMessages.append("FAILED: PartitionedStorageUnit's stored data was negative.")
+        partIndex = list(dataNameVec[-1]).index(partName)
+        dataIndex = list(partitionName).index(partName)
+        if partName == "test":
+            baudRate = 0.0
+        elif partName == "node_1_msg":
+            baudRate = baudRate_1
+        elif partName == "node_2_msg":
+            baudRate = baudRate_2
 
-    if testFailCount:
-        print(testMessages)
-    else:
-        print("Passed")
-    return [testFailCount, ''.join(testMessages)]
+        # Check 2 - if partition exists, does it added data to the correct partition?
+        np.testing.assert_allclose(dataVec[-1][partIndex], 
+            add_data[dataIndex] + baudRate*sim_time, 
+            err_msg = (
+            "FAILED: PartitionedStorageUnit did not use the correct partition."))
 
 
 if __name__ == "__main__":
-    print(test_module(False))
\ No newline at end of file
+    baudRate_1 = 1200
+    baudRate_2 = 1200
+    storageCapacity = 2400
+    expectedStorage = 2400
+    test_storage_limits(baudRate_1, baudRate_2, storageCapacity, expectedStorage)
+    add_data = [1200, 200]
+    partitionName = ["test", "node_1_msg"]
+    test_set_data_buffer(baudRate_1, baudRate_2, partitionName, 
+                         add_data, storageCapacity, expectedStorage)
+    storageCapacity = 20000
+    test_set_data_buffer_partition(baudRate_1, baudRate_2, partitionName, 
+                         add_data, storageCapacity)
\ No newline at end of file
diff --git a/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_simpleStorageUnit.py b/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_simpleStorageUnit.py
index 6d3cdecfa4..8dc3da100b 100644
--- a/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_simpleStorageUnit.py
+++ b/src/simulation/onboardDataHandling/storageUnit/_UnitTest/test_simpleStorageUnit.py
@@ -17,6 +17,8 @@
 
 import inspect
 import os
+import pytest
+import numpy as np
 
 filename = inspect.getframeinfo(inspect.currentframe()).filename
 path = os.path.dirname(os.path.abspath(filename))
@@ -25,30 +27,28 @@
 
 # Import all of the modules that we are going to be called in this simulation
 from Basilisk.utilities import SimulationBaseClass
-from Basilisk.utilities import unitTestSupport                  # general support file with common unit test functions
 from Basilisk.simulation import simpleStorageUnit
 from Basilisk.architecture import messaging
 from Basilisk.utilities import macros
 
-def test_module(show_plots):
-    # each test method requires a single assert method to be called
-    [testResults, testMessage] = checkStorage_limits(show_plots)
-    assert testResults < 1, testMessage
+params_storage_limits = [(1200, 1200, 2400, 2400),
+                     (600, 1200, 3600, 3600),
+                     (600, 600, 10000, 6000)]
 
+@pytest.mark.parametrize("baudRate_1, baudRate_2, storageCapacity, expectedStorage",
+                          params_storage_limits)
 
-def checkStorage_limits(show_plots):
+def test_storage_limits(baudRate_1, baudRate_2, storageCapacity, expectedStorage):
     """
     Tests:
 
-    1. Whether the simpleStorageUnit can add multiple nodes (core base class functionality);
-    2. That the simpleStorageUnit correctly evaluates how much stored data it should have given a pair of
-       1200 baud input messages.
+    1. Whether the simpleStorageUnit can add multiple nodes (core base class 
+        functionality);
+    2. That the simpleStorageUnit correctly evaluates how much stored data it should 
+        have given a pair of input messages.
 
-    :param show_plots: Not used; no plots to be shown.
     """
 
-    testFailCount = 0                       # zero unit test result counter
-    testMessages = []                       # create empty array to store test log messages
     unitTaskName = "unitTask"               # arbitrary name (don't change)
     unitProcessName = "TestProcess"         # arbitrary name (don't change)
 
@@ -61,15 +61,15 @@ def checkStorage_limits(show_plots):
     testProc.addTask(unitTestSim.CreateNewTask(unitTaskName, testProcessRate))
 
     test_storage_unit = simpleStorageUnit.SimpleStorageUnit()
-    test_storage_unit.storageCapacity = 2400. # bit capacity.
+    test_storage_unit.storageCapacity = storageCapacity # bit capacity.
 
     dataMsg1 = messaging.DataNodeUsageMsgPayload()
-    dataMsg1.baudRate = 1200. # baud
+    dataMsg1.baudRate = baudRate_1 # baud
     dataMsg1.dataName = "node_1_msg"
     dat1Msg = messaging.DataNodeUsageMsg().write(dataMsg1)
 
     dataMsg2 = messaging.DataNodeUsageMsgPayload()
-    dataMsg2.baudRate = 1200. # baud
+    dataMsg2.baudRate = baudRate_2 # baud
     dataMsg2.dataName = "node_2_msg"
     dat2Msg = messaging.DataNodeUsageMsg().write(dataMsg2)
 
@@ -90,38 +90,150 @@ def checkStorage_limits(show_plots):
     storedDataLog = dataLog.storageLevel
     capacityLog = dataLog.storageCapacity
     netBaudLog = dataLog.currentNetBaud
+    partitionName = dataLog.storedDataName
 
-    #   Check 1 - is net baud rate equal to 2400.?
+    #   Check 1 - is net baud rate correct?
     for ind in range(0,len(netBaudLog)):
         currentBaud = netBaudLog[ind]
-        if currentBaud != 2400.:
-            testFailCount += 1
-            testMessages.append("FAILED: SimpleStorageUnit did not correctly log the net baud rate.")
+        np.testing.assert_allclose(currentBaud, baudRate_1 + baudRate_2, atol=1e-1,
+            err_msg=("FAILED: PartitionedStorageUnit did not correctly log baud rate."))
 
-    #print(netBaudLog)
+    #   Check 2 - is used storage space correct?
+    np.testing.assert_allclose(storedDataLog[-1], expectedStorage, atol=1e-4,
+        err_msg=("FAILED: PartitionedStorageUnit did not track integrated data."))
 
-    if not unitTestSupport.isDoubleEqualRelative((2400.),storedDataLog[-1], 1e-8):
-        testFailCount+=1
-        testMessages.append("FAILED: SimpleStorageUnit did not track integrated data. Returned "+str(storedDataLog[-1,1])+", expected "+str((2400.)))
+    #   Check 3 - is the amount of data more than zero and less than the capacity?
+    for ind in range(0,len(storedDataLog)):
+        assert storedDataLog[ind] <= capacityLog[ind] or np.isclose(storedDataLog[ind], 
+            capacityLog[ind]), (
+            "FAILED: PartitionedStorageUnit's stored data exceeded its capacity.")
+
+        assert storedDataLog[ind] >= 0., (
+            "FAILED: PartitionedStorageUnit's stored data was negative.")
+
+    #   Check 4 - is there only one partition?
+    assert len(partitionName[0]) == 1, (
+        "FAILED: PartitionedStorageUnit did use the correct partition.")
+    
+    #   Check 6 - is the name of the partition correct?
+    assert partitionName[0][0] == "STORED DATA", (
+        "FAILED: PartitionedStorageUnit did not correctly log the partition name.")
+
+
+params_set_data = [(1200, 1200, 1200, 2400, 2400),
+                   (600, 600, 0, 10000, 6000),
+                   (600, 600, 4e3, 10000, 10000),
+                   (0, 0, 1000, 2000, 1000),
+                   (0, 0, -1000, 2000, 0),
+                   (1000, 0, -2000, 6e3, 3000),
+                   (0, 0, 3000, 2000, 0),
+                   (600, 0, 3000, 10000, 6000),
+                   (600, 600, 3000, 10000, 9000),
+                   (300, 600, 3000, 10000, 7500)]
+
+@pytest.mark.parametrize(
+        "baudRate_1, baudRate_2, add_data, storageCapacity, expectedStorage",
+        params_set_data)
+def test_set_data_buffer(baudRate_1, baudRate_2, add_data, storageCapacity, 
+                         expectedStorage):
+    """
+    Tests:
 
-    #print(storedDataLog)
+    1. Whether the partitionedStorageUnit can add data using the setDataBuffer method;
+    2. That the partitionedStorageUnit correctly evaluates how much stored data it 
+        should have given a pair of input messages and using setDataBuffer.
 
-    for ind in range(0,len(storedDataLog)):
-        if storedDataLog[ind] > capacityLog[ind]:
-            testFailCount +=1
-            testMessages.append("FAILED: SimpleStorageUnit's stored data exceeded its capacity.")
+    :return:
+    """
+
+    unitTaskName = "unitTask"               # arbitrary name (don't change)
+    unitProcessName = "TestProcess"         # arbitrary name (don't change)
+
+    # Create a sim module as an empty container
+    unitTestSim = SimulationBaseClass.SimBaseClass()
+
+    # Create test thread
+    testProcessRate = macros.sec2nano(0.1)     # update process rate update time
+    testProc = unitTestSim.CreateNewProcess(unitProcessName)
+    testProc.addTask(unitTestSim.CreateNewTask(unitTaskName, testProcessRate))
+
+    test_storage_unit = simpleStorageUnit.SimpleStorageUnit()
+    test_storage_unit.storageCapacity = storageCapacity # bit capacity.
 
-        if storedDataLog[ind] < 0.:
-            testFailCount +=1
-            testMessages.append("FAILED: SimpleStorageUnit's stored data was negative.")
+    dataMsg1 = messaging.DataNodeUsageMsgPayload()
+    dataMsg1.baudRate = baudRate_1 # baud
+    dataMsg1.dataName = "node_1_msg"
+    dat1Msg = messaging.DataNodeUsageMsg().write(dataMsg1)
 
-    if testFailCount:
-        print(testMessages)
-    else:
-        print("Passed")
+    dataMsg2 = messaging.DataNodeUsageMsgPayload()
+    dataMsg2.baudRate = baudRate_2 # baud
+    dataMsg2.dataName = "node_2_msg"
+    dat2Msg = messaging.DataNodeUsageMsg().write(dataMsg2)
 
-    return [testFailCount, ''.join(testMessages)]
+    # Test the addNodeToStorage method:
+    test_storage_unit.addDataNodeToModel(dat1Msg)
+    test_storage_unit.addDataNodeToModel(dat2Msg)
+
+    unitTestSim.AddModelToTask(unitTaskName, test_storage_unit)
+
+    dataLog = test_storage_unit.storageUnitDataOutMsg.recorder()
+    unitTestSim.AddModelToTask(unitTaskName, dataLog)
+    test_storage_unit.setDataBuffer(0)
+
+    unitTestSim.InitializeSimulation()
+    sim_time = 5.0
+    unitTestSim.ConfigureStopTime(macros.sec2nano(sim_time - 1.0))
+    unitTestSim.ExecuteSimulation()
+
+    test_storage_unit.setDataBuffer(add_data)
+    unitTestSim.ConfigureStopTime(macros.sec2nano(sim_time))
+    unitTestSim.ExecuteSimulation()
+
+    storedDataLog = dataLog.storageLevel
+    capacityLog = dataLog.storageCapacity
+    netBaudLog = dataLog.currentNetBaud
+    partitionName = dataLog.storedDataName
+    partitionData = dataLog.storedData
+
+    #   Check 1 - is net baud rate correct?
+    for ind in range(0,len(netBaudLog)):
+        currentBaud = netBaudLog[ind]
+        np.testing.assert_allclose(currentBaud, baudRate_1 + baudRate_2, atol=1e-4,
+            err_msg=("FAILED: PartitionedStorageUnit did not correctly log baud rate."))
+
+    #   Check 2 - is used storage space correct?
+    np.testing.assert_allclose(storedDataLog[-1], expectedStorage, atol=1e-4,
+        err_msg=("FAILED: PartitionedStorageUnit did not track integrated data."))
+
+    #   Check 3 - is the amount of data more than zero and less than the capacity?
+    for ind in range(0,len(storedDataLog)):
+        assert storedDataLog[ind] <= capacityLog[ind] or np.isclose(storedDataLog[ind], 
+            capacityLog[ind]), (
+            "FAILED: PartitionedStorageUnit's stored data exceeded its capacity.")
+
+        assert storedDataLog[ind] >= 0., (
+            "FAILED: PartitionedStorageUnit's stored data was negative.")
+
+    #   Check 4 - is the data in the partitioned storage unit correct?
+    assert partitionData[-1][0] == storedDataLog[-1], (
+        "FAILED: PartitionedStorageUnit did not correctly log the stored data.")
+    
+    #   Check 5 - is there only one partition?
+    assert len(partitionName[0]) == 1, (
+        "FAILED: PartitionedStorageUnit should have just one partition.")
+   
+    #   Check 6 - is the name of the partition correct?
+    assert partitionName[0][0] == "STORED DATA", (
+        "FAILED: PartitionedStorageUnit did not correctly log the partition name.")
 
 
 if __name__ == "__main__":
-    print(checkStorage_limits(False))
\ No newline at end of file
+    baudRate_1 = 1200
+    baudRate_2 = 1200
+    storageCapacity = 2400
+    expectedStorage = 2400
+    test_storage_limits(baudRate_1, baudRate_2, 
+                        storageCapacity, expectedStorage)
+    add_data = 1200
+    test_set_data_buffer(baudRate_1, baudRate_2, add_data, 
+                         storageCapacity, expectedStorage)
\ No newline at end of file

From 9b5dc30fecae62948641aeeb93e2b1eb540a23db Mon Sep 17 00:00:00 2001
From: LorenzzoQM <lorenzzo.quevedomantovani@colorado.edu>
Date: Wed, 22 Nov 2023 11:15:19 -0700
Subject: [PATCH 4/4] Added description of setDataBuffer to .rst files and
 ReleaseNotes

---
 docs/source/Support/bskReleaseNotes.rst                       | 1 +
 .../_GeneralModuleFiles/dataStorageUnitBase.rst               | 3 ++-
 .../storageUnit/partitionedStorageUnit.rst                    | 4 ++++
 .../onboardDataHandling/storageUnit/simpleStorageUnit.rst     | 4 ++++
 4 files changed, 11 insertions(+), 1 deletion(-)

diff --git a/docs/source/Support/bskReleaseNotes.rst b/docs/source/Support/bskReleaseNotes.rst
index 94ff382d86..d480d060c8 100644
--- a/docs/source/Support/bskReleaseNotes.rst
+++ b/docs/source/Support/bskReleaseNotes.rst
@@ -82,6 +82,7 @@ Version |release|
   use spherical harmonics and loads them from a file with a single command. Similarly, the methods ``usePolyhedralGravityModel``
   and ``usePointMassGravityModel`` have been added.
 - Fixed examples and tests to run even when Basilisk is built with ``--vizInterface False``.
+- Added a new method ``setDataBuffer()`` to :ref:`simpleStorageUnit` and :ref:`partitionedStorageUnit` to add or remove data from specified partitions.
 
 Version 2.2.0 (June 28, 2023)
 -----------------------------
diff --git a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.rst b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.rst
index 0090ffd709..036f8cbdbe 100644
--- a/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.rst
+++ b/src/simulation/onboardDataHandling/_GeneralModuleFiles/dataStorageUnitBase.rst
@@ -2,10 +2,11 @@ Executive Summary
 -----------------
 DataStorageUnitBase is a base class that is used generate a standard interface and list of features for modules that store simulated onboard data.  This class is used by other modules as a parent class and cannot be instantiated by itself.  All Basilisk data storage modules based on this DataStorageUnitBase inherit the following common properties:
 
-1. Writes out a :ref:`DataStorageStatusMsgPayload` containing the sum of the current stored data (in bits), the storage capacity (bits), the current net data rate (in baud), an array of char array containing the names of the stored data (ex. Instrument 1, Instrument 2), and an array of doubles containing the stored data associated with each type (bits).
+1. Writes out a :ref:`DataStorageStatusMsgPayload` containing the sum of the current stored data (in bits), the storage capacity (bits), the current net data rate (in baud), an array of char array containing the names of the stored data (ex. Instrument 1, Instrument 2), and an array of integers containing the stored data associated with each type (bits).
 2. Allows for multiple :ref:`DataNodeUsageMsgPayload` corresponding to individual :ref:`dataNodeBase` instances to be subscribed to using the ``addDataNodeToModel(msg)`` method.
 3. Iterates through attached :ref:`DataNodeUsageMsgPayload` instances, integrates the data for each data node, and adds it to its respective entry using ``integrateDataStatus()`` method, which may be overwritten in child classes.
 4. Loops through the vector of storedData to sum the total amount of data contained within the storage unit.
+5. Add or remove data from specified partitions using ``setDataBuffer()`` method.
 
 Core functionality is wrapped in the ``integrateDataStatus`` protected virtual void method, which computes the amount of data stored in a storage unit on a module basis. This base class automatically implements a partitioned storage unit (different data buffers for each device). See :ref:`simpleStorageUnit` for an example of how this functionality can be overwritten.
 
diff --git a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.rst b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.rst
index bc7761001b..24eb16c17a 100644
--- a/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.rst
+++ b/src/simulation/onboardDataHandling/storageUnit/partitionedStorageUnit.rst
@@ -40,4 +40,8 @@ Then, the names of the partitions need to be added to the storageUnit using::
 
    storageUnit.addPartition("partitionName")
 
+The ``setDataBuffer()`` method can be used to add or remove a given amount of data from specified partitions::
+
+   storageUnit.setDataBuffer(["partitionName","anotherPartitionName"], [1E4, -1E4]) # Given in bits
+
 For more information on how to set up and use this module, see the simple data system example :ref:`scenarioDataDemo`.
diff --git a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.rst b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.rst
index 897c1f6db9..f149ac6f43 100644
--- a/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.rst
+++ b/src/simulation/onboardDataHandling/storageUnit/simpleStorageUnit.rst
@@ -35,4 +35,8 @@ The next step is to attach one or more :ref:`DataNodeUsageMsgPayload` instances
 
    storageUnit.addDataNodeToModel(dataMsg)
 
+The method ``setDataBuffer()`` can be used to add or remove a specific amount of data from the storage unit::
+
+   storageUnit.setDataBuffer(1E4) # Given in bits
+
 For more information on how to set up and use this module, see the simple data system example :ref:`scenarioDataDemo`.