2 files changed, 5 insertions, 5 deletions
diff --git a/tests/federation/test_federation.py b/tests/federation/test_federation.py
index c893c4863d..01f07ff36d 100644
--- a/tests/federation/test_federation.py
+++ b/tests/federation/test_federation.py
@@ -257,7 +257,7 @@ class FederationTestCase(unittest.TestCase):
         response = yield self.federation.make_query(
             destination="remote",
             query_type="a-question",
-            args={"one": "1", "two": "2"}
+            args={"one": "1", "two": "2"},
         )
 
         self.assertEquals({"your": "response"}, response)
@@ -265,7 +265,8 @@ class FederationTestCase(unittest.TestCase):
         self.mock_http_client.get_json.assert_called_with(
             destination="remote",
             path="/_matrix/federation/v1/query/a-question",
-            args={"one": "1", "two": "2"}
+            args={"one": "1", "two": "2"},
+            retry_on_dns_fail=True,
         )
 
     @defer.inlineCallbacks
diff --git a/tests/handlers/test_directory.py b/tests/handlers/test_directory.py
index dd5d85dde6..e10a49a8ac 100644
--- a/tests/handlers/test_directory.py
+++ b/tests/handlers/test_directory.py
@@ -20,7 +20,6 @@ from twisted.internet import defer
 from mock import Mock
 
 from synapse.server import HomeServer
-from synapse.http.client import HttpClient
 from synapse.handlers.directory import DirectoryHandler
 from synapse.storage.directory import RoomAliasMapping
 
@@ -95,8 +94,8 @@ class DirectoryTestCase(unittest.TestCase):
             query_type="directory",
             args={
                 "room_alias": "#another:remote",
-                HttpClient.RETRY_DNS_LOOKUP_FAILURES: False
-            }
+            },
+            retry_on_dns_fail=False,
         )
 
     @defer.inlineCallbacks