Custom TestRail API wrapper

trail.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import urllib2
import json
import base64


class TestRailAPIError(Exception):

    def __init__(self, code, error):
        self.code = code
        self.error = error

    def __str__(self):
        return '[%s] %s' % (self.code, self.error)


def url_join(*parts):
    return '/'.join([part.strip('/') for part in parts])


class APIClient(object):

    def __init__(self, base_url, login='', token=''):

        auth = base64.b64encode('%s:%s' % (login, token))
        self.__headers = {
            'Authorization': 'Basic %s' % auth,
            'Content-Type': 'application/json; charset=utf-8'
        }

        self.__url = url_join(base_url, 'index.php?', 'api', 'v2')

    def get(self, path):
        return self.__send_request(path)

    def post(self, path, data):
        return self.__send_request(path, data)

    def __send_request(self, path, data=None):
        url = url_join(self.__url, path)

        data_str = json.dumps(data) if data else None
        request = urllib2.Request(url, headers=self.__headers, data=data_str)

        try:
            response = urllib2.urlopen(request).read()
        except urllib2.HTTPError as e:
            code = e.code
            error = json.loads(e.read()).get('error', '<err msg missed>')

            raise TestRailAPIError(code, error)

        return json.loads(response) if response else {}


class CustomAPIClient(APIClient):

    def __init__(self, base_url, login='', token=''):
        super(CustomAPIClient, self).__init__(base_url, login, token)

    def get_status_mapping(self):
        statuses = self.get('/get_statuses')
        return {status['id']: status['label'] for status in statuses}

    def get_users_mapping(self):
        users = self.get('/get_users')
        return {user['id']: user['email'] for user in users}

    def get_projects(self):
        projects = self.get('/get_projects')
        return {project['id']: project['name'] for project in projects}

    def get_milestones(self, project_id):
        milestones = self.get('/get_milestones/%s' % project_id)
        return {mstone['id']: mstone['name'] for mstone in milestones}

    def get_plans(self, project_id, milestone_id):
        plans = self.get('/get_plans/%s&milestone_id=%s' % (project_id, milestone_id))
        return plans

    def get_plan(self, plan_id):
        return self.get('/get_plan/%s' % plan_id)

    def get_counts_for_plan(self, plan_id):
        plan_details = self.get_plan(plan_id)
        retest = plan_details.get('retest_count', 0)
        passed = plan_details.get('passed_count', 0)
        failed = plan_details.get('failed_count', 0)
        blocked = plan_details.get('blocked_count', 0)
        untested = plan_details.get('untested_count', 0)

        total = sum((retest, passed, failed, blocked, untested))

        return {
            'retest': retest,
            'passed': passed,
            'failed': failed,
            'blocked': blocked,
            'untested': untested,
            'total': total
        }

    def get_tests_for_run(self, run_id):
        return self.get('/get_tests/%s' % run_id)

    def get_tests_for_plan(self, plan_id):
        run_ids = list()
        for entry in self.get_plan(plan_id)['entries']:
            run_ids.extend([i['id'] for i in entry['runs']])

        test_cases = list()
        for run_id in list(set(run_ids)):
            run_cases = self.get_tests_for_run(run_id)
            test_cases.extend(run_cases)

        return test_cases