diff --git a/src/energy/examples/basic-energy-model-test.cc b/src/energy/examples/basic-energy-model-test.cc
index fe2957113..7abfe39e5 100644
--- a/src/energy/examples/basic-energy-model-test.cc
+++ b/src/energy/examples/basic-energy-model-test.cc
@@ -187,7 +187,6 @@ BasicEnergyUpdateTest::StateSwitchTest (WifiPhy::State state)
   double voltage = source->GetSupplyVoltage ();
   estRemainingEnergy -= devModel->GetIdleCurrentA () * voltage * m_timeS;
 
-
   // calculate new state power consumption
   double current = 0.0;
   switch (state)
@@ -210,11 +209,15 @@ BasicEnergyUpdateTest::StateSwitchTest (WifiPhy::State state)
     case WifiPhy::SLEEP:
       current = devModel->GetSleepCurrentA ();
       break;
+    case WifiPhy::OFF:
+      current = 0;
+      break;
     default:
       NS_FATAL_ERROR ("Undefined radio state: " << state);
       break;
     }
   estRemainingEnergy -= current * voltage * m_timeS;
+  estRemainingEnergy = std::max(0.0, estRemainingEnergy);
 
   // obtain remaining energy from source
   double remainingEnergy = source->GetRemainingEnergy ();
@@ -354,7 +357,7 @@ BasicEnergyDepletionTest::DepletionTestCase (double simTimeS,
   WifiHelper wifi;
   wifi.SetStandard (WIFI_PHY_STANDARD_80211b);
 
-  YansWifiPhyHelper wifiPhy =  YansWifiPhyHelper::Default ();
+  YansWifiPhyHelper wifiPhy = YansWifiPhyHelper::Default ();
   /*
    * This is one parameter that matters when using FixedRssLossModel, set it to
    * zero; otherwise, gain will be added.
diff --git a/src/energy/test/examples-to-run.py b/src/energy/test/examples-to-run.py
index c1f97e86c..1cba3c38d 100644
--- a/src/energy/test/examples-to-run.py
+++ b/src/energy/test/examples-to-run.py
@@ -10,6 +10,7 @@
 cpp_examples = [
     ("li-ion-energy-source", "True", "True"),
     ("rv-battery-model-test", "True", "True"),
+    ("basic-energy-model-test", "True", "True"),
 ]
 
 # A list of Python examples to run in order to ensure that they remain