Merge remote branch 'cros/upstream' into master.

client/tests/wb_kupdate/wb_kupdate_unittest.py

Index: client/tests/wb_kupdate/wb_kupdate_unittest.py
diff --git a/client/tests/wb_kupdate/wb_kupdate_unittest.py b/client/tests/wb_kupdate/wb_kupdate_unittest.py
new file mode 100644
index 0000000000000000000000000000000000000000..84092ed699c9e253112ac4348111330bd77471ce
--- /dev/null
+++ b/client/tests/wb_kupdate/wb_kupdate_unittest.py
@@ -0,0 +1,105 @@
+#!/usr/bin/python
+
+import common
+import datetime
+import logging
+import os
+import time
+import unittest
+from autotest_lib.client.bin import test
+from autotest_lib.client.bin import utils
+from autotest_lib.client.common_lib import error
+from autotest_lib.client.common_lib.test_utils import mock
+from autotest_lib.client.tests.wb_kupdate import wb_kupdate
+
+class WbKupdateUnitTest(unittest.TestCase):
+    def setUp(self):
+        """Set up all required variables for the Unittest.
+        """
+        self._logger = logging.getLogger()
+        self._wbkupdate_obj = WbKupdateSubclass()
+        self._god = mock.mock_god()
+
+    def test_needs_more_time(self):
+        """Tests the _needs_more_time method.
+        """
+        self._logger.info('Testing the "_needs_more_time" method.')
+
+        # Obvious failure - since start_time < start_time + 1.
+        self.assertTrue(self._wbkupdate_obj._needs_more_time(
+                start_time=datetime.datetime.now(),
+                duration=1))
+
+        # Check if 1 minute has elapsed since start_time.
+        self.assertFalse(self._wbkupdate_obj._needs_more_time(
+                start_time=datetime.datetime.now(),
+                duration=1,
+                _now=datetime.datetime.now() + datetime.timedelta(seconds=60)))
+
+    def test_wait_until_data_flushed_pass(self):
+        """Tests the _wait_until_data_flushed method.
+
+        This tests the "success" code path.
+        """
+        self._logger.info('Testing the "_wait_until_data_flushed" method - '
+                          'Success code path.')
+
+        # Creating stubs for required methods.
+        self._god.stub_function(self._wbkupdate_obj,
+                                "_get_disk_usage")
+
+        # Setting default return values for stub functions.
+        # Setting the initial size of the file.
+        self._wbkupdate_obj._get_disk_usage.expect_call('').and_return(10)
+        # Returning the same file size - forcing code path to enter loop.
+        self._wbkupdate_obj._get_disk_usage.expect_call('').and_return(10)
+        # Returning a greater file size - exiting the while loop.
+        self._wbkupdate_obj._get_disk_usage.expect_call('').and_return(11)
+
+        # Call the method.
+        self._wbkupdate_obj._wait_until_data_flushed(datetime.datetime.now(),
+                                                     1)
+
+        # Ensure all stubbed methods called.
+        self._god.check_playback()
+
+
+    def test_wait_until_data_flushed_fail(self):
+        """Tests the _wait_until_data_flushed method.
+
+        This tests the "failure" code path.
+        """
+        self._logger.info('Testing the "_wait_until_data_flushed" method - '
+                          'Failure code path.')
+        # Creating stubs for required methods.
+        self._god.stub_function(self._wbkupdate_obj,
+                                "_get_disk_usage")
+
+        # Setting default return values for stub functions.
+        # Setting the initial size of the file.
+        self._wbkupdate_obj._get_disk_usage.expect_call('').and_return(10)
+        # Returning the same file size - forcing code path to enter loop.
+        self._wbkupdate_obj._get_disk_usage.expect_call('').and_return(10)
+
+        # Call the method.
+        self.assertRaises(error.TestError,
+                          self._wbkupdate_obj._wait_until_data_flushed,
+                          start_time=datetime.datetime.now(),
+                          max_wait_time=0)
+
+        # Ensure all stubbed methods called.
+        self._god.check_playback()
+
+
+class WbKupdateSubclass(wb_kupdate.wb_kupdate):
+    """Sub-classing the wb_kupdate class.
+    """
+    def __init__(self):
+        """Empty constructor.
+        """
+        # Create all test defaults.
+        self.initialize()
+
+
+if __name__ == '__main__':
+    unittest.main()