1 files changed, 46 insertions, 2 deletions
diff --git a/tests/pa/test_provision.py b/tests/pa/test_provision.py
index a98eacd..0446442 100644
--- a/tests/pa/test_provision.py
+++ b/tests/pa/test_provision.py
@@ -16,17 +16,22 @@
 #
 
 import azurelinuxagent.common.utils.fileutil as fileutil
+
+from azurelinuxagent.common.exception import ProtocolError
 from azurelinuxagent.common.osutil.default import DefaultOSUtil
 from azurelinuxagent.common.protocol import OVF_FILE_NAME
 from azurelinuxagent.pa.provision import get_provision_handler
+from azurelinuxagent.pa.provision.default import ProvisionHandler
 from tests.tools import *
 
 
 class TestProvision(AgentTestCase):
  
     @distros("redhat")
-    def test_provision(self, distro_name, distro_version, distro_full_name):
-        provision_handler = get_provision_handler(distro_name, distro_version, 
+    @patch('azurelinuxagent.common.osutil.default.DefaultOSUtil.get_instance_id',
+        return_value='B9F3C233-9913-9F42-8EB3-BA656DF32502')
+    def test_provision(self, mock_util, distro_name, distro_version, distro_full_name):
+        provision_handler = get_provision_handler(distro_name, distro_version,
                                                   distro_full_name)
         mock_osutil = MagicMock()
         mock_osutil.decode_customdata = Mock(return_value="")
@@ -48,6 +53,45 @@ class TestProvision(AgentTestCase):
         data = DefaultOSUtil().decode_customdata(base64data)
         fileutil.write_file(tempfile.mktemp(), data)
 
+    @patch('os.path.isfile', return_value=False)
+    def test_is_provisioned_not_provisioned(self, mock_isfile):
+        ph = ProvisionHandler()
+        self.assertFalse(ph.is_provisioned())
+
+    @patch('os.path.isfile', return_value=True)
+    @patch('azurelinuxagent.common.utils.fileutil.read_file',
+            return_value="B9F3C233-9913-9F42-8EB3-BA656DF32502")
+    @patch('azurelinuxagent.pa.deprovision.get_deprovision_handler')
+    def test_is_provisioned_is_provisioned(self,
+            mock_deprovision, mock_read, mock_isfile):
+        ph = ProvisionHandler()
+        ph.osutil = Mock()
+        ph.osutil.get_instance_id = \
+            Mock(return_value="B9F3C233-9913-9F42-8EB3-BA656DF32502")
+        ph.write_provisioned = Mock()
+
+        deprovision_handler = Mock()
+        mock_deprovision.return_value = deprovision_handler
+
+        self.assertTrue(ph.is_provisioned())
+        deprovision_handler.run_changed_unique_id.assert_not_called()
+
+    @patch('os.path.isfile', return_value=True)
+    @patch('azurelinuxagent.common.utils.fileutil.read_file',
+            side_effect=["Value"])
+    @patch('azurelinuxagent.pa.deprovision.get_deprovision_handler')
+    def test_is_provisioned_not_deprovisioned(self,
+            mock_deprovision, mock_read, mock_isfile):
+
+        ph = ProvisionHandler()
+        ph.osutil = Mock()
+        ph.write_provisioned = Mock()
+
+        deprovision_handler = Mock()
+        mock_deprovision.return_value = deprovision_handler
+
+        self.assertTrue(ph.is_provisioned())
+        deprovision_handler.run_changed_unique_id.assert_called_once()
 
 if __name__ == '__main__':
     unittest.main()