From a5ab01390fa289915f2b3dbcf57076ac4b0445a8 Mon Sep 17 00:00:00 2001
From: Guillaume COTER <guillaume.coter@gmail.com>
Date: Sat, 26 Sep 2020 22:09:00 +0000
Subject: [PATCH] remote[ssh]: Add the parameter allow_agent (#4623)

This parameter is passed to paramiko.
See http://docs.paramiko.org/en/stable/api/client.html?highlight=allow_agent#paramiko.client.SSHClient.connect for more details.

Fixes #4186
---
 dvc/config.py                     |  1 +
 dvc/tree/ssh/__init__.py          |  2 ++
 tests/unit/remote/ssh/test_ssh.py | 25 +++++++++++++++++++++++++
 3 files changed, 28 insertions(+)

diff --git a/dvc/config.py b/dvc/config.py
index a01e39291..372a51d1f 100644
--- a/dvc/config.py
+++ b/dvc/config.py
@@ -190,6 +190,7 @@ SCHEMA = {
                     "keyfile": str,
                     "timeout": Coerce(int),
                     "gss_auth": Bool,
+                    "allow_agent": Bool,
                     **REMOTE_COMMON,
                 },
                 "hdfs": {"user": str, **REMOTE_COMMON},
diff --git a/dvc/tree/ssh/__init__.py b/dvc/tree/ssh/__init__.py
index 2faa6ec5f..d6ba6a89a 100644
--- a/dvc/tree/ssh/__init__.py
+++ b/dvc/tree/ssh/__init__.py
@@ -91,6 +91,7 @@ class SSHTree(BaseTree):
             self.sock = paramiko.ProxyCommand(proxy_command)
         else:
             self.sock = None
+        self.allow_agent = config.get("allow_agent", True)
 
     @staticmethod
     def ssh_config_filename():
@@ -143,6 +144,7 @@ class SSHTree(BaseTree):
             password=self.password,
             gss_auth=self.gss_auth,
             sock=self.sock,
+            allow_agent=self.allow_agent,
         )
 
     @contextmanager
diff --git a/tests/unit/remote/ssh/test_ssh.py b/tests/unit/remote/ssh/test_ssh.py
index be6cc704f..7c15b00f8 100644
--- a/tests/unit/remote/ssh/test_ssh.py
+++ b/tests/unit/remote/ssh/test_ssh.py
@@ -183,6 +183,31 @@ def test_ssh_gss_auth(mock_file, mock_exists, dvc, config, expected_gss_auth):
     assert tree.gss_auth == expected_gss_auth
 
 
+@pytest.mark.parametrize(
+    "config,expected_allow_agent",
+    [
+        ({"url": "ssh://example.com"}, True),
+        ({"url": "ssh://not_in_ssh_config.com"}, True),
+        ({"url": "ssh://example.com", "allow_agent": True}, True),
+        ({"url": "ssh://example.com", "allow_agent": False}, False),
+    ],
+)
+@patch("os.path.exists", return_value=True)
+@patch(
+    f"{builtin_module_name}.open",
+    new_callable=mock_open,
+    read_data=mock_ssh_config,
+)
+def test_ssh_allow_agent(
+    mock_file, mock_exists, dvc, config, expected_allow_agent
+):
+    tree = SSHTree(dvc, config)
+
+    mock_exists.assert_called_with(SSHTree.ssh_config_filename())
+    mock_file.assert_called_with(SSHTree.ssh_config_filename())
+    assert tree.allow_agent == expected_allow_agent
+
+
 def test_hardlink_optimization(dvc, tmp_dir, ssh):
     tree = SSHTree(dvc, ssh.config)