unit tests, updated custom header getters/setters

Author: SamHausmann

examples/language.py

@@ -16,15 +16,9 @@ def run(key, altUrl='https://api.rosette.com/rest/v1/'):
     api = API(user_key=key, service_url=altUrl)
 
     language_data = "Por favor Señorita, says the man."
-    app_header = []
-    app_header.append("X-RosetteAPI-App")
-    app_header.append("python-app")
-    headers = []
-    headers.append(app_header)
     params = DocumentParameters()
-
     params["content"] = language_data
-    params["customHeaders"] = headers
+    api.setCustomHeaders("X-RosetteAPI-App", "python-app")
     return api.language(params)
 
 

rosette/api.py

@@ -26,6 +26,7 @@
 import os
 from socket import gaierror
 import requests
+import re
 
 _BINDING_VERSION = '1.1.1'
 _GZIP_BYTEARRAY = bytearray([0x1F, 0x8b, 0x08])
@@ -186,7 +187,7 @@ class DocumentParameters(_DocumentParamSetBase):
     def __init__(self):
         """Create a L{DocumentParameters} object."""
         _DocumentParamSetBase.__init__(
-            self, ("content", "contentUri", "language", "genre", "customHeaders"))
+            self, ("content", "contentUri", "language", "genre"))
         self.file_name = ""
         self.useMultipart = False
 
@@ -268,8 +269,7 @@ def __init__(self):
              "sourceScript",
              "targetScript",
              "targetScheme",
-             "genre",
-             "customHeaders"))
+             "genre"))
 
     def validate(self):
         """Internal. Do not use."""
@@ -302,7 +302,7 @@ class NameSimilarityParameters(_DocumentParamSetBase):
 
     def __init__(self):
         self.useMultipart = False
-        _DocumentParamSetBase.__init__(self, ("name1", "name2", "customHeaders"))
+        _DocumentParamSetBase.__init__(self, ("name1", "name2"))
 
     def validate(self):
         """Internal. Do not use."""
@@ -367,6 +367,17 @@ def info(self):
         identifying data."""
         url = self.service_url + "info"
         headers = {'Accept': 'application/json', 'X-RosetteAPI-Binding': 'python', 'X-RosetteAPI-Binding-Version': _BINDING_VERSION}
+
+        customHeaders = self.api.getCustomHeaders()
+        pattern = re.compile('^X-RosetteAPI-')
+        if customHeaders is not None:
+            for key in customHeaders.keys():
+                if pattern.match(key) is not None:
+                    headers[key] = customHeaders[key]
+                else:
+                    raise RosetteException("badHeader", "Custom header name must begin with \"X-RosetteAPI-\"", key);
+            self.api.clearCustomHeaders()
+
         if self.debug:
             headers['X-RosetteAPI-Devel'] = 'true'
         self.logger.info('info: ' + url)
@@ -383,6 +394,17 @@ def ping(self):
 
         url = self.service_url + 'ping'
         headers = {'Accept': 'application/json', 'X-RosetteAPI-Binding': 'python', 'X-RosetteAPI-Binding-Version': _BINDING_VERSION}
+
+        customHeaders = self.api.getCustomHeaders()
+        pattern = re.compile('^X-RosetteAPI-')
+        if customHeaders is not None:
+            for key in customHeaders.keys():
+                if pattern.match(key) is not None:
+                    headers[key] = customHeaders[key]
+                else:
+                    raise RosetteException("badHeader", "Custom header name must begin with \"X-RosetteAPI-\"", key);
+            self.api.clearCustomHeaders()
+
         if self.debug:
             headers['X-RosetteAPI-Devel'] = 'true'
         self.logger.info('Ping: ' + url)
@@ -428,14 +450,20 @@ def call(self, parameters):
         headers = {}
         if self.user_key is not None:
 
-            if parameters["customHeaders"] is not None:
-                for h in parameters["customHeaders"]:
-                    headers[h[0]] = h[1]
-                parameters["customHeaders"] = None
+            customHeaders = self.api.getCustomHeaders()
+            pattern = re.compile('^X-RosetteAPI-')
+            if customHeaders is not None:
+                for key in customHeaders.keys():
+                    if pattern.match(key) is not None:
+                        headers[key] = customHeaders[key]
+                    else:
+                        raise RosetteException("badHeader", "Custom header name must begin with \"X-RosetteAPI-\"", key);
+                self.api.clearCustomHeaders()
 
             headers["X-RosetteAPI-Key"] = self.user_key
             headers["X-RosetteAPI-Binding"] = "python"
             headers["X-RosetteAPI-Binding-Version"] = _BINDING_VERSION
+
         if self.useMultipart:
             params = dict(
                 (key,
@@ -509,6 +537,7 @@ def __init__(
         self.connection_refresh_duration = refresh_duration
         self.http_connection = None
         self.options = {}
+        self.customHeaders = {}
 
     def _connect(self, parsedUrl):
         """ Simple connection method
@@ -645,6 +674,30 @@ def clearOptions(self):
         """
         self.options.clear()
 
+    def setCustomHeaders(self, name, value):
+        """
+        Sets custom headers
+
+        @param headers: array of custom headers to be set
+        """
+        if value is None:
+            self.customHeaders.pop(name, None)
+        else:
+            self.customHeaders[name] = value
+
+    def getCustomHeaders(self):
+        """
+        Get custom headers
+        """
+        return self.customHeaders
+
+    def clearCustomHeaders(self):
+        """
+        Clears custom headers
+        """
+
+        self.customHeaders.clear()
+
     def ping(self):
         """
         Create a ping L{EndpointCaller} for the server and ping it.

tests/test_rosette_api.py

@@ -82,10 +82,35 @@ def test_option_clear_single_option(api):
     api.setOption('test', None)
     assert api.getOption('test') is None
 
+# Test the custom header set/get/clear
+
+
+def test_custom_header_get_set_clear(api):
+    key = 'X-RosetteAPI-Test'
+    value = 'foo'
+    api.setCustomHeaders(key, value)
+    assert value == api.getCustomHeaders()[key]
+
+    api.clearCustomHeaders()
+    assert len(api.getCustomHeaders()) is 0
+
+# Test for invalid header name
+
+
+def test_invalid_header(api):
+    key = 'test'
+    value = 'foo'
+    api.setCustomHeaders(key, value)
+
+    with pytest.raises(RosetteException) as e_rosette:
+        result = api.info()
+
+    assert e_rosette.value.status == 'badHeader'
 
 # Test that pinging the API is working properly
 # @httpretty.activate
 
+
 def test_ping(api, json_response):
     httpretty.enable()
     httpretty.register_uri(httpretty.GET, "https://api.rosette.com/rest/v1/ping",