Changed hostmap to be a named argument for compatibility

Author: zephraph

browsermobproxy/client.py

@@ -281,18 +281,18 @@ def timeouts(self, options):
                          params)
         return r.status_code
 
-    def remap_hosts(self, hostmap):
+    def remap_hosts(self, address=None, ip_address=None, hostmap={}):
         """
         Remap the hosts for a specific URL
 
-        :param hostmap: A dictionary of url/ip_address pairs to remap
+        :param address: url that you wish to remap
+        :param ip_address: IP Address that will handle all traffic for the address passed in
+        :param **hostmap: Other hosts to be added as keyword arguments
         """
-        if not isinstance(hostmap, dict):
-            raise TypeError("hostmap needs to be a dictionary")
-  
-        if not hostmap:
-            raise Exception("hostmap should have at least one entry")
-  
+
+        if (address is not None and ip_address is not None):
+            hostmap[address] = ip_address
+
         r = requests.post('%s/proxy/%s/hosts' % (self.host, self.port),
                          json.dumps(hostmap),
                           headers={'content-type': 'application/json'})

test/test_client.py

@@ -225,7 +225,14 @@ def test_remap_hosts(self):
         """
         /proxy/:port/hosts
         """
-        status_code = self.client.remap_hosts({"example.com": "1.2.3.4"})
+        status_code = self.client.remap_hosts("example.com", "1.2.3.4")
+        assert(status_code == 200)
+
+    def test_remap_hosts_with_hostmap(self):
+        """
+        /proxy/:port/hosts
+        """
+        status_code = self.client.remap_hosts(hostmap={"example.com": "1.2.3.4"})
         assert(status_code == 200)
 
     def test_wait_for_traffic_to_stop(self):