Code cleanup/refactoring (#31)

This is cleanup based on flake8 and pylint recommendations + corrections to make implementation more generic (for future extensibility of support POST/PUT/DELETE calls)

Update: also added new-application API to fix #30

Author: dimon222

tests/test_base.py

@@ -24,7 +24,7 @@ def test_valid_request(self):
             requests_get_mock.get('/ololo', text=json.dumps(BaseYarnAPITestCase.success_response()))
 
             client = self.get_client()
-            response = client.request('/ololo', foo='bar')
+            response = client.request('/ololo', params={"foo": 'bar'})
 
             assert requests_get_mock.called
             self.assertIn(response.data['status'], 'success')
@@ -34,7 +34,7 @@ def test_valid_request_with_parameters(self):
             requests_get_mock.get('/ololo?foo=bar', text=json.dumps(BaseYarnAPITestCase.success_response()))
 
             client = self.get_client()
-            response = client.request('/ololo', foo='bar')
+            response = client.request('/ololo', params={"foo": 'bar'})
 
             assert requests_get_mock.called
             self.assertIn(response.data['status'], 'success')

tests/test_hadoop_conf.py

@@ -165,6 +165,24 @@ def test_get_jobhistory_host_port(self):
             host_port = hadoop_conf.get_jobhistory_host_port()
             self.assertIsNone(host_port)
 
+
+    def test_get_nodemanager_host_port(self):
+        with patch('yarn_api_client.hadoop_conf.parse') as parse_mock:
+            parse_mock.return_value = 'example.com:8022'
+
+            host_port = hadoop_conf.get_nodemanager_host_port()
+
+            self.assertEqual(('example.com', '8022'), host_port)
+            parse_mock.assert_called_with('/etc/hadoop/conf/yarn-site.xml',
+                                          'yarn.nodemanager.webapp.address')
+
+            parse_mock.reset_mock()
+            parse_mock.return_value = None
+
+            host_port = hadoop_conf.get_nodemanager_host_port()
+            self.assertIsNone(host_port)
+
+
     def test_get_webproxy_host_port(self):
         with patch('yarn_api_client.hadoop_conf.parse') as parse_mock:
             parse_mock.return_value = 'example.com:8022'

tests/test_history_server.py

@@ -23,17 +23,21 @@ def test_application_information(self, request_mock):
 
     def test_jobs(self, request_mock):
         self.hs.jobs()
-        request_mock.assert_called_with('/ws/v1/history/mapreduce/jobs')
+        request_mock.assert_called_with('/ws/v1/history/mapreduce/jobs', params={})
 
         self.hs.jobs(state='NEW', user='root', queue='high', limit=100,
                      started_time_begin=1, started_time_end=2,
                      finished_time_begin=3, finished_time_end=4)
 
         request_mock.assert_called_with('/ws/v1/history/mapreduce/jobs',
-                                        queue='high',
-                                        state='NEW', user='root', limit=100,
-                                        startedTimeBegin=1, startedTimeEnd=2,
-                                        finishedTimeBegin=3, finishedTimeEnd=4)
+                                        params={"queue": 'high',
+                                                "state": 'NEW',
+                                                "user": 'root',
+                                                "limit": 100,
+                                                "startedTimeBegin": 1,
+                                                "startedTimeEnd": 2,
+                                                "finishedTimeBegin": 3,
+                                                "finishedTimeEnd": 4})
 
         with self.assertRaises(IllegalArgumentError):
             self.hs.jobs(state='ololo')
@@ -56,12 +60,12 @@ def test_job_conf(self, request_mock):
 
     def test_job_tasks(self, request_mock):
         self.hs.job_tasks('job_2')
-        request_mock.assert_called_with('/ws/v1/history/mapreduce/jobs/job_2/tasks')
-        self.hs.job_tasks('job_2', type='m')
-        request_mock.assert_called_with('/ws/v1/history/mapreduce/jobs/job_2/tasks', type='m')
+        request_mock.assert_called_with('/ws/v1/history/mapreduce/jobs/job_2/tasks', params={})
+        self.hs.job_tasks('job_2', job_type='m')
+        request_mock.assert_called_with('/ws/v1/history/mapreduce/jobs/job_2/tasks', params={"type": 'm'})
 
         with self.assertRaises(IllegalArgumentError):
-            self.hs.job_tasks('job_2', type='ololo')
+            self.hs.job_tasks('job_2', job_type='ololo')
 
     def test_job_task(self, request_mock):
         self.hs.job_task('job_2', 'task_3')

tests/test_node_manager.py

@@ -18,10 +18,10 @@ def test_node_information(self, request_mock):
     def test_node_applications(self, request_mock):
         self.nm.node_applications('RUNNING', 'root')
         request_mock.assert_called_with('/ws/v1/node/apps',
-                                        state='RUNNING', user='root')
+                                        params={"state":'RUNNING', "user":'root'})
 
         self.nm.node_applications()
-        request_mock.assert_called_with('/ws/v1/node/apps')
+        request_mock.assert_called_with('/ws/v1/node/apps', params={})
 
         with self.assertRaises(IllegalArgumentError):
             self.nm.node_applications('ololo', 'root')

tests/test_resource_manager.py

@@ -64,15 +64,40 @@ def test_cluster_application_attempts(self, request_mock):
 
     def test_cluster_nodes(self, request_mock):
         self.rm.cluster_nodes()
-        request_mock.assert_called_with('/ws/v1/cluster/nodes')
+        request_mock.assert_called_with('/ws/v1/cluster/nodes', params={})
 
         self.rm.cluster_nodes(state='NEW', healthy='true')
         request_mock.assert_called_with('/ws/v1/cluster/nodes',
-                                        state='NEW', healthy='true')
+                                        params={"state": 'NEW', "healthy": 'true'})
 
         with self.assertRaises(IllegalArgumentError):
             self.rm.cluster_nodes(state='NEW', healthy='ololo')
 
     def test_cluster_node(self, request_mock):
         self.rm.cluster_node('node_1')
         request_mock.assert_called_with('/ws/v1/cluster/nodes/node_1')
+
+    # TODO
+    # def test_cluster_submit_application(self, request_mock):
+    #     self.rm.cluster_submit_application()
+    #     request_mock.assert_called_with('/ws/v1/cluster/apps')
+
+    def test_cluster_new_application(self, request_mock):
+        self.rm.cluster_new_application()
+        request_mock.assert_called_with('/ws/v1/cluster/apps/new-application', 'POST')
+
+    def test_cluster_get_application_queue(self, request_mock):
+        self.rm.cluster_get_application_queue('app_1')
+        request_mock.assert_called_with('/ws/v1/cluster/apps/app_1/queue')
+
+    def test_cluster_change_application_queue(self, request_mock):
+        self.rm.cluster_change_application_queue('app_1', 'queue_1')
+        request_mock.assert_called_with('/ws/v1/cluster/apps/app_1/queue', 'PUT', data={"queue": 'queue_1'})
+
+    def test_cluster_get_application_priority(self, request_mock):
+        self.rm.cluster_get_application_priority('app_1')
+        request_mock.assert_called_with('/ws/v1/cluster/apps/app_1/priority')
+
+    def test_cluster_change_application_priority(self, request_mock):
+        self.rm.cluster_change_application_priority('app_1', 'priority_1')
+        request_mock.assert_called_with('/ws/v1/cluster/apps/app_1/priority', 'PUT', data={"priority": 'priority_1'})

yarn_api_client/__init__.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 __version__ = '0.3.4.dev0'
-__all__ = ['ApplicationMaster', 'HistoryServer', 'NodeManager','ResourceManager']
+__all__ = ['ApplicationMaster', 'HistoryServer', 'NodeManager', 'ResourceManager']
 
 from .application_master import ApplicationMaster
 from .history_server import HistoryServer

yarn_api_client/application_master.py

@@ -21,11 +21,11 @@ class ApplicationMaster(BaseYarnAPI):
     :param boolean kerberos_enabled: Flag identifying is Kerberos Security has been enabled for YARN
     """
     def __init__(self, address=None, port=8088, timeout=30, kerberos_enabled=False):
-        self.address, self.port, self.timeout, self.kerberos_enabled = address, port, timeout, kerberos_enabled
         if address is None:
             self.logger.debug('Get configuration from hadoop conf dir')
             address, port = get_webproxy_host_port()
-            self.address, self.port = address, port
+
+        super(ApplicationMaster, self).__init__(address, port, timeout, kerberos_enabled)
 
     def application_information(self, application_id):
         """

yarn_api_client/base.py

@@ -1,7 +1,6 @@
 # -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 
-import json
 import logging
 import requests
 
@@ -17,46 +16,33 @@ class BaseYarnAPI(object):
     __logger = None
     response_class = Response
 
+    def __init__(self, address=None, port=None, timeout=None, kerberos_enabled=None):
+        self.address, self.port, self.timeout, self.kerberos_enabled = address, port, timeout, kerberos_enabled
+
     def _validate_configuration(self):
         if self.address is None:
             raise ConfigurationError('API address is not set')
         elif self.port is None:
             raise ConfigurationError('API port is not set')
 
-    def request(self, api_path, **query_args):
-        params = query_args
+    def request(self, api_path, method='GET', **kwargs):
         api_endpoint = 'http://{}:{}{}'.format(self.address, self.port, api_path)
 
         self.logger.info('API Endpoint {}'.format(api_endpoint))
 
         self._validate_configuration()
 
-        response = None
-        if self.kerberos_enabled:
-            from requests_kerberos import HTTPKerberosAuth
-            response = requests.get(api_endpoint, params, auth=HTTPKerberosAuth())
-        else:
-            response = requests.get(api_endpoint, params)
-
-        if response.status_code == requests.codes.ok:
-            return self.response_class(response)
+        if method == 'GET':
+            headers = None
         else:
-            msg = 'Response finished with status: %s. Details: %s' % (response.status_code, response.text)
-            raise APIError(msg)
-
-    def update(self, api_path, data):
-        api_endpoint = 'http://{}:{}{}'.format(self.address, self.port, api_path)
-
-        self.logger.info('API Endpoint {}'.format(api_endpoint))
-
-        self._validate_configuration()
+            headers = {"Content-Type": "application/json"}
 
         response = None
         if self.kerberos_enabled:
             from requests_kerberos import HTTPKerberosAuth
-            response = requests.put(api_endpoint, data=data, auth=HTTPKerberosAuth(), headers={"Content-Type":"application/json"})
+            response = requests.request(method=method, url=api_endpoint, auth=HTTPKerberosAuth(), headers=headers, **kwargs)
         else:
-            response = requests.put(api_endpoint, data=data, headers={"Content-Type":"application/json"})
+            response = requests.request(method=method, url=api_endpoint, headers=headers, **kwargs)
 
         if response.status_code == requests.codes.ok:
             return self.response_class(response)

yarn_api_client/hadoop_conf.py

@@ -16,7 +16,7 @@ def _get_rm_ids(hadoop_conf_path):
     return rm_ids
 
 
-def _get_resource_manager(hadoop_conf_path, rm_id = None):
+def _get_resource_manager(hadoop_conf_path, rm_id=None):
     prop_name = 'yarn.resourcemanager.webapp.address'
     if rm_id is not None:
         rm_webapp_address = parse(os.path.join(hadoop_conf_path, 'yarn-site.xml'), '%s.%s' % (prop_name, rm_id))
@@ -70,6 +70,17 @@ def get_jobhistory_host_port():
         return None
 
 
+def get_nodemanager_host_port():
+    config_path = os.path.join(CONF_DIR, 'yarn-site.xml')
+    prop_name = 'yarn.nodemanager.webapp.address'
+    value = parse(config_path, prop_name)
+    if value is not None:
+        host, _, port = value.partition(':')
+        return host, port
+    else:
+        return None
+
+
 def get_webproxy_host_port():
     config_path = os.path.join(CONF_DIR, 'yarn-site.xml')
     prop_name = 'yarn.web-proxy.address'

yarn_api_client/history_server.py

@@ -21,11 +21,11 @@ class HistoryServer(BaseYarnAPI):
     :param boolean kerberos_enabled: Flag identifying is Kerberos Security has been enabled for YARN
     """
     def __init__(self, address=None, port=19888, timeout=30, kerberos_enabled=False):
-        self.address, self.port, self.timeout, self.kerberos_enabled = address, port, timeout, kerberos_enabled
         if address is None:
             self.logger.debug('Get information from hadoop conf dir')
             address, port = get_jobhistory_host_port()
-            self.address, self.port = address, port
+
+        super(HistoryServer, self).__init__(address, port, timeout, kerberos_enabled)
 
     def application_information(self):
         """
@@ -65,7 +65,7 @@ def jobs(self, state=None, user=None, queue=None, limit=None,
         """
         path = '/ws/v1/history/mapreduce/jobs'
 
-        legal_states = set([s for s, _ in JobStateInternal])
+        legal_states = {s for s, _ in JobStateInternal}
         if state is not None and state not in legal_states:
             msg = 'Job Internal State %s is illegal' % (state,)
             raise IllegalArgumentError(msg)
@@ -82,7 +82,7 @@ def jobs(self, state=None, user=None, queue=None, limit=None,
 
         params = self.construct_parameters(loc_args)
 
-        return self.request(path, **params)
+        return self.request(path, params=params)
 
     def job(self, job_id):
         """
@@ -134,7 +134,7 @@ def job_conf(self, job_id):
 
         return self.request(path)
 
-    def job_tasks(self, job_id, type=None):
+    def job_tasks(self, job_id, job_type=None):
         """
         With the tasks API, you can obtain a collection of resources that
         represent a task within a job.
@@ -151,15 +151,15 @@ def job_tasks(self, job_id, type=None):
         # m - for map
         # r - for reduce
         valid_types = ['m', 'r']
-        if type is not None and type not in valid_types:
-            msg = 'Job type %s is illegal' % (type,)
+        if job_type is not None and job_type not in valid_types:
+            msg = 'Job type %s is illegal' % (job_type,)
             raise IllegalArgumentError(msg)
 
         params = {}
-        if type is not None:
-            params['type'] = type
+        if job_type is not None:
+            params['type'] = job_type
 
-        return self.request(path, **params)
+        return self.request(path, params=params)
 
     def job_task(self, job_id, task_id):
         """