2 files changed, 169 insertions, 18 deletions
diff --git a/tests/common/osutil/test_default.py b/tests/common/osutil/test_default.py
index d982b7e..933787b 100644
--- a/tests/common/osutil/test_default.py
+++ b/tests/common/osutil/test_default.py
@@ -21,6 +21,7 @@ import mock
 import azurelinuxagent.common.osutil.default as osutil
 import azurelinuxagent.common.utils.shellutil as shellutil
 from azurelinuxagent.common.osutil import get_osutil
+from azurelinuxagent.common.utils import fileutil
 from tests.tools import *
 
 
@@ -164,5 +165,155 @@ class TestOSUtil(AgentTestCase):
         else:
             self.fail("Cannot retrieve number of process cores using shell command.")
 
+    def test_conf_sshd(self):
+        new_file = "\
+Port 22\n\
+Protocol 2\n\
+ChallengeResponseAuthentication yes\n\
+#PasswordAuthentication yes\n\
+UsePAM yes\n\
+"
+        expected_output = "\
+Port 22\n\
+Protocol 2\n\
+ChallengeResponseAuthentication no\n\
+#PasswordAuthentication yes\n\
+UsePAM yes\n\
+PasswordAuthentication no\n\
+ClientAliveInterval 180\n\
+"
+
+        with patch.object(fileutil, 'write_file') as patch_write:
+            with patch.object(fileutil, 'read_file', return_value=new_file):
+                osutil.DefaultOSUtil().conf_sshd(disable_password=True)
+                patch_write.assert_called_once_with(
+                    conf.get_sshd_conf_file_path(),
+                    expected_output)
+
+    def test_conf_sshd_with_match(self):
+        new_file = "\
+Port 22\n\
+ChallengeResponseAuthentication yes\n\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+"
+        expected_output = "\
+Port 22\n\
+ChallengeResponseAuthentication no\n\
+PasswordAuthentication no\n\
+ClientAliveInterval 180\n\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+"
+
+        with patch.object(fileutil, 'write_file') as patch_write:
+            with patch.object(fileutil, 'read_file', return_value=new_file):
+                osutil.DefaultOSUtil().conf_sshd(disable_password=True)
+                patch_write.assert_called_once_with(
+                    conf.get_sshd_conf_file_path(),
+                    expected_output)
+
+    def test_conf_sshd_with_match_last(self):
+        new_file = "\
+Port 22\n\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+"
+        expected_output = "\
+Port 22\n\
+PasswordAuthentication no\n\
+ChallengeResponseAuthentication no\n\
+ClientAliveInterval 180\n\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+"
+
+        with patch.object(fileutil, 'write_file') as patch_write:
+            with patch.object(fileutil, 'read_file', return_value=new_file):
+                osutil.DefaultOSUtil().conf_sshd(disable_password=True)
+                patch_write.assert_called_once_with(
+                    conf.get_sshd_conf_file_path(),
+                    expected_output)
+
+    def test_conf_sshd_with_match_middle(self):
+        new_file = "\
+Port 22\n\
+match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+match all\n\
+#Other config\n\
+"
+        expected_output = "\
+Port 22\n\
+match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+match all\n\
+#Other config\n\
+PasswordAuthentication no\n\
+ChallengeResponseAuthentication no\n\
+ClientAliveInterval 180\n\
+"
+
+        with patch.object(fileutil, 'write_file') as patch_write:
+            with patch.object(fileutil, 'read_file', return_value=new_file):
+                osutil.DefaultOSUtil().conf_sshd(disable_password=True)
+                patch_write.assert_called_once_with(
+                    conf.get_sshd_conf_file_path(),
+                    expected_output)
+
+    def test_conf_sshd_with_match_multiple(self):
+        new_file = "\
+Port 22\n\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+Match host 192.168.1.2\n\
+  ChallengeResponseAuthentication yes\n\
+Match all\n\
+#Other config\n\
+"
+        expected_output = "\
+Port 22\n\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+Match host 192.168.1.2\n\
+  ChallengeResponseAuthentication yes\n\
+Match all\n\
+#Other config\n\
+PasswordAuthentication no\n\
+ChallengeResponseAuthentication no\n\
+ClientAliveInterval 180\n\
+"
+
+        with patch.object(fileutil, 'write_file') as patch_write:
+            with patch.object(fileutil, 'read_file', return_value=new_file):
+                osutil.DefaultOSUtil().conf_sshd(disable_password=True)
+                patch_write.assert_called_once_with(
+                    conf.get_sshd_conf_file_path(),
+                    expected_output)
+
+    def test_conf_sshd_with_match_multiple_first_last(self):
+        new_file = "\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+Match host 192.168.1.2\n\
+  ChallengeResponseAuthentication yes\n\
+"
+        expected_output = "\
+PasswordAuthentication no\n\
+ChallengeResponseAuthentication no\n\
+ClientAliveInterval 180\n\
+Match host 192.168.1.1\n\
+  ChallengeResponseAuthentication yes\n\
+Match host 192.168.1.2\n\
+  ChallengeResponseAuthentication yes\n\
+"
+
+        with patch.object(fileutil, 'write_file') as patch_write:
+            with patch.object(fileutil, 'read_file', return_value=new_file):
+                osutil.DefaultOSUtil().conf_sshd(disable_password=True)
+                patch_write.assert_called_once_with(
+                    conf.get_sshd_conf_file_path(),
+                    expected_output)
+
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/common/test_version.py b/tests/common/test_version.py
index 306ea16..df4f0e3 100644
--- a/tests/common/test_version.py
+++ b/tests/common/test_version.py
@@ -17,29 +17,13 @@
 
 from __future__ import print_function
 
-import copy
-import glob
-import json
-import mock
-import os
-import platform
-import random
-import subprocess
-import sys
-import tempfile
 import textwrap
-import zipfile
 
-from tests.protocol.mockwiredata import *
-from tests.tools import *
+import mock
 
-import azurelinuxagent.common.conf as conf
-import azurelinuxagent.common.logger as logger
-import azurelinuxagent.common.utils.fileutil as fileutil
 import azurelinuxagent.common.version as version
-
-from azurelinuxagent.common.utils.flexible_version import FlexibleVersion
 from azurelinuxagent.common.version import *
+from tests.tools import *
 
 
 class TestCurrentAgentName(AgentTestCase):
@@ -54,6 +38,22 @@ class TestCurrentAgentName(AgentTestCase):
         self.assertEqual(AGENT_VERSION, str(current_version))
         return
 
+    @patch("os.getcwd", return_value="/")
+    def test_extract_name_root_finds_installed(self, mock_cwd):
+        current_agent, current_version = set_current_agent()
+        self.assertEqual(AGENT_LONG_VERSION, current_agent)
+        self.assertEqual(AGENT_VERSION, str(current_version))
+        return
+
+    @patch("os.getcwd")
+    def test_extract_name_in_path_finds_installed(self, mock_cwd):
+        path = os.path.join(conf.get_lib_dir(), "events")
+        mock_cwd.return_value = path
+        current_agent, current_version = set_current_agent()
+        self.assertEqual(AGENT_LONG_VERSION, current_agent)
+        self.assertEqual(AGENT_VERSION, str(current_version))
+        return
+
     @patch("os.getcwd")
     def test_extract_name_finds_latest_agent(self, mock_cwd):
         path = os.path.join(conf.get_lib_dir(), "{0}-{1}".format(