--- a/src/netanim/test/netanim-test.cc	Tue Dec 10 09:12:38 2013 +0400
+++ b/src/netanim/test/netanim-test.cc	Tue Dec 10 09:12:38 2013 +0400
@@ -26,55 +26,120 @@
 #include "ns3/netanim-module.h"
 #include "ns3/applications-module.h"
 #include "ns3/point-to-point-layout-module.h"
+#include "ns3/basic-energy-source.h"
+#include "ns3/simple-device-energy-model.h"
 
 namespace ns3 {
 
-class AnimationInterfaceTestCase : public TestCase
+class AbstractAnimationInterfaceTestCase : public TestCase
+{
+public:
+  /**
+   * \brief Constructor.
+   */
+  AbstractAnimationInterfaceTestCase(std::string name);
+  /**
+   * \brief Destructor.
+   */
+  virtual
+  ~AbstractAnimationInterfaceTestCase();
+  /**
+   * \brief Run unit tests for this class.
+   * \return false if all tests have passed, false otherwise
+   */
+  virtual void
+  DoRun(void);
+
+protected:
+
+  NodeContainer m_nodes;
+  AnimationInterface* m_anim;
+
+private:
+
+  virtual void
+  PrepareNetwork() = 0;
+
+  virtual void
+  CheckLogic() = 0;
+
+  virtual void
+  CheckFileExistence();
+
+  const char* m_traceFileName;
+};
+
+AbstractAnimationInterfaceTestCase::AbstractAnimationInterfaceTestCase (std::string name) :
+  TestCase (name), m_anim(NULL), m_traceFileName("netanim-test.xml")
+{
+}
+
+AbstractAnimationInterfaceTestCase::~AbstractAnimationInterfaceTestCase ()
+{
+  delete m_anim;
+}
+
+void
+AbstractAnimationInterfaceTestCase::DoRun (void)
+{
+  PrepareNetwork();
+
+  m_anim = new AnimationInterface(m_traceFileName);
+
+  Simulator::Run ();
+  CheckLogic();
+  CheckFileExistence();
+  Simulator::Destroy ();
+}
+
+void
+AbstractAnimationInterfaceTestCase::CheckFileExistence()
+{
+  FILE * fp = fopen (m_traceFileName, "r");
+  NS_TEST_ASSERT_MSG_NE (fp, 0, "Trace file was not created");
+  fclose (fp);
+  unlink (m_traceFileName);
+}
+
+class AnimationInterfaceTestCase : public AbstractAnimationInterfaceTestCase
 {
 public:
   /**
    * \brief Constructor.
    */
   AnimationInterfaceTestCase ();
-  /**
-   * \brief Destructor.
-   */
-  virtual
-  ~AnimationInterfaceTestCase ();
-  /**
-   * \brief Run unit tests for this class.
-   * \return false if all tests have passed, false otherwise
-   */
+
+private:
+
   virtual void
-  DoRun (void);
+  PrepareNetwork();
+
+  virtual void
+  CheckLogic();
 
 };
 
 AnimationInterfaceTestCase::AnimationInterfaceTestCase () :
-  TestCase ("Verify AnimationInterface")
+  AbstractAnimationInterfaceTestCase ("Verify AnimationInterface")
 {
 }
 
-AnimationInterfaceTestCase::~AnimationInterfaceTestCase ()
+void
+AnimationInterfaceTestCase::PrepareNetwork (void)
 {
-}
-void
-AnimationInterfaceTestCase::DoRun (void)
-{
-  NodeContainer nodes;
-  nodes.Create (2);
-  AnimationInterface::SetConstantPosition (nodes.Get (0), 0 , 10);
-  AnimationInterface::SetConstantPosition (nodes.Get (1), 1 , 10);
+  m_nodes.Create (2);
+  AnimationInterface::SetConstantPosition (m_nodes.Get (0), 0 , 10);
+  AnimationInterface::SetConstantPosition (m_nodes.Get (1), 1 , 10);
 
   PointToPointHelper pointToPoint;
   pointToPoint.SetDeviceAttribute ("DataRate", StringValue ("5Mbps"));
   pointToPoint.SetChannelAttribute ("Delay", StringValue ("2ms"));
 
   NetDeviceContainer devices;
-  devices = pointToPoint.Install (nodes);
+  devices = pointToPoint.Install (m_nodes);
 
   InternetStackHelper stack;
-  stack.Install (nodes);
+  stack.Install (m_nodes);
 
   Ipv4AddressHelper address;
   address.SetBase ("10.1.1.0", "255.255.255.0");
@@ -83,7 +148,7 @@
 
   UdpEchoServerHelper echoServer (9);
 
-  ApplicationContainer serverApps = echoServer.Install (nodes.Get (1));
+  ApplicationContainer serverApps = echoServer.Install (m_nodes.Get (1));
   serverApps.Start (Seconds (1.0));
   serverApps.Stop (Seconds (10.0));
 
@@ -92,18 +157,78 @@
   echoClient.SetAttribute ("Interval", TimeValue (Seconds (1.0)));
   echoClient.SetAttribute ("PacketSize", UintegerValue (1024));
 
-  ApplicationContainer clientApps = echoClient.Install (nodes.Get (0));
+  ApplicationContainer clientApps = echoClient.Install (m_nodes.Get (0));
   clientApps.Start (Seconds (2.0));
   clientApps.Stop (Seconds (10.0));
-  std::string traceFileName = "netanim-test.xml";
-  AnimationInterface anim(traceFileName.c_str ());
-  Simulator::Run ();
-  NS_TEST_ASSERT_MSG_EQ (anim.GetTracePktCount (), 32, "Expected 32 packets traced");
-  FILE * fp = fopen (traceFileName.c_str (), "r");
-  NS_TEST_ASSERT_MSG_NE (fp, 0, "Trace file was not created");
-  fclose (fp);
-  unlink (traceFileName.c_str ());
-  Simulator::Destroy ();
+}
+
+void
+AnimationInterfaceTestCase::CheckLogic (void)
+{
+  NS_TEST_ASSERT_MSG_EQ (m_anim->GetTracePktCount (), 32, "Expected 32 packets traced");
+}
+
+class AnimationRemainingEnergyTestCase : public AbstractAnimationInterfaceTestCase
+{
+public:
+  /**
+   * \brief Constructor.
+   */
+  AnimationRemainingEnergyTestCase ();
+
+private:
+
+  virtual void
+  PrepareNetwork();
+
+  virtual void
+  CheckLogic();
+
+  Ptr<BasicEnergySource> m_energySource;
+  Ptr<SimpleDeviceEnergyModel> m_energyModel;
+  ObjectFactory m_energySourceFactory;
+  ObjectFactory m_deviceEnergyModelFactory;
+  const double m_initialEnergy;
+};
+
+AnimationRemainingEnergyTestCase::AnimationRemainingEnergyTestCase () :
+  AbstractAnimationInterfaceTestCase("Verify Remaining energy tracing"),
+  m_initialEnergy(100)
+{
+  m_energySourceFactory.SetTypeId ("ns3::BasicEnergySource");
+  m_deviceEnergyModelFactory.SetTypeId ("ns3::SimpleDeviceEnergyModel");
+}
+
+void
+AnimationRemainingEnergyTestCase::PrepareNetwork (void)
+{
+  m_energySource = m_energySourceFactory.Create<BasicEnergySource>();
+  m_energyModel = m_deviceEnergyModelFactory.Create<SimpleDeviceEnergyModel>();
+
+  m_energySource->SetInitialEnergy(m_initialEnergy);
+  m_energyModel->SetEnergySource (m_energySource);
+  m_energySource->AppendDeviceEnergyModel (m_energyModel);
+  m_energyModel->SetCurrentA(20);
+
+  m_nodes.Create (1);
+  AnimationInterface::SetConstantPosition (m_nodes.Get (0), 0 , 10);
+
+  // aggregate energy source to node
+  m_nodes.Get (0)->AggregateObject (m_energySource);
+  // once node's energy will be depleted according to the model
+  Simulator::Stop (Seconds (2));
+}
+
+void
+AnimationRemainingEnergyTestCase::CheckLogic (void)
+{
+  const double remainingEnergy = m_energySource->GetRemainingEnergy();
+
+  NS_TEST_ASSERT_MSG_EQ((remainingEnergy < m_initialEnergy), true, "Energy hasn't depleted!");
+  NS_TEST_ASSERT_MSG_EQ_TOL(m_anim->GetNodeEnergyFraction(m_nodes.Get (0)),
+                            remainingEnergy / m_initialEnergy,
+                            1.0e-13,
+                            "Wrong remaining energy value was traced");
 }
 
 static class AnimationInterfaceTestSuite : public TestSuite
@@ -113,6 +238,7 @@
     TestSuite ("animation-interface", UNIT)
   {
     AddTestCase (new AnimationInterfaceTestCase (), TestCase::QUICK);
+    AddTestCase (new AnimationRemainingEnergyTestCase (), TestCase::QUICK);
   }
 } g_animationInterfaceTestSuite;