Skip to content
Snippets Groups Projects
py3.py 1.69 KiB
Newer Older
import json
import requests
from urllib3 import Retry

class HttpClient(object):
    """
    Local client for batch-running and testing purposes
    """

    app_name = 'model_server.base.api:app'

    def __init__(self, host='127.0.0.1', port=6221, raise_nonsuccess=False, **kwargs) -> None:
        self.uri = f'http://{host}:{port}/'
        self.raise_nonsuccess = raise_nonsuccess

    def _get_sesh(self):
        sesh = requests.Session()
        retries = Retry(
            total=5,
            backoff_factor=0.1,
        )
        sesh.mount('http://', requests.adapters.HTTPAdapter(max_retries=retries))
        return sesh

    def get(self, endpoint):
        resp = self._get_sesh().get(self.uri + endpoint)
        if self.raise_nonsuccess and resp.status_code != 200:
            raise NonSuccessResposeError(resp.text)
        return resp

    def put(self, endpoint, query=None, body=None,):
        resp = self._get_sesh().put(
            self.uri + endpoint,
            params=query,
            data=json.dumps(body)
        if self.raise_nonsuccess and resp.status_code != 200:
            raise NonSuccessResponseError(resp.text)
        return resp

    def hit(self, method, endpoint, params=None, body=None):
        """
        Automatically issue a request based on method, mainly for parsing JSON configurations
        """
        if method.upper() == 'GET':
            return self.get(endpoint)
        elif method.upper() == 'PUT':
            return self.put(endpoint, params, body)
        else:
            raise InvalidMethodError(f'Invalid method {method}')

class Error(Exception):
    pass

class InvalidMethodError(Error):
    pass

class NonSuccessResponseError(Error):