Source code for schrodinger.test.performance.client

import json

import requests
from requests.models import Response

import schrodinger.test.performance.exceptions as exceptions
import schrodinger.test.stu.client as stu_client
from schrodinger.test.performance import reporter
from schrodinger.test.stu import common


class PerformanceTestClient(stu_client._BaseClient):

    def __init__(self, username=None):
        super().__init__(username=username)
        self._productapi = self.item_uri('product')
        self._testapi = self.performance_uri('test')

    def performance_uri(self, item):
        return f"{self._base}/performance{self._api}{item}/"

    def getSingleEntityFromResponse(
        self,
        response: Response,
        exception_missing_key_message:
        str = "Response does not contain the 'objects' key",
        exception_mulitple_objects_msg:
        str = "Response contains multiple 'objects', expecting only one"
    ) -> dict:
        """
        Take a response object, extract JSON, validate its contents.
        Return a single object if exists, return None of zero objects found, raise BedResponse exception otherwise.
        """
        data = response.json()
        if 'objects' not in data:
            raise exceptions.BadResponse(exception_missing_key_message)
        objects = data['objects']
        if len(objects) == 1:
            return objects[0]
        elif len(objects) > 1:
            raise exceptions.BadResponse(exception_mulitple_objects_msg)
        else:
            return None

    def get_or_create(self, url: str, params: dict) -> str:
        """Get or create a resource matching the parameters."""
        response = self.get(url, params=params)
        entity = self.getSingleEntityFromResponse(response)
        if entity:
            return entity['resource_uri']
        response = self.post(
            url,
            data=json.dumps(params),
        )
        location = response.headers['location']
        return location.replace(common.BASE_URL, '')

    def get_or_create_test(self,
                           name: str,
                           description: str,
                           product_name: str,
                           scival=False):
        """
        Get or create a single test from the performance database.

        Setting `scival` to `True` will add the 'scival' tag when creating a new test.
        """
        params = dict(name=product_name)
        response = self.get(self._productapi,
                            params=params,
                            required_statuses=(200, 401, 404))
        username = stu_client.get_stu_username()
        no_product_msg = (
            f'No product named "{product_name}". See the list of product names '
            f'at {common.BASE_URL}/products. File a BLDMGR JIRA ticket if you '
            'need to add a product to STU. Please specify the name of the '
            'the product, the TL associated with the product, and the jira '
            'project which tracks tickets as part of the project.')
        if response.status_code == 404:
            raise exceptions.BadResponse(no_product_msg)
        try:
            response.raise_for_status()
        except requests.exceptions.HTTPError as http_error:
            if response.status_code == 401:
                raise exceptions.BadResponse(
                    f"{http_error}, please verify that the appropriate user is making this request: username='{username}'"
                )
            raise
        entity = self.getSingleEntityFromResponse(
            response, exception_missing_key_message=no_product_msg)
        if not entity:
            raise exceptions.BadResponse(no_product_msg)

        product = entity['resource_uri']
        product_id = reporter.resource_id(product)

        # Get an existing test:
        test_dict = dict(name=name, product=product_id)
        response = self.get(self._testapi, params=test_dict)
        entity = self.getSingleEntityFromResponse(response)
        if entity:
            return entity['resource_uri']

        # Create a new test:
        if not description:
            raise ValueError(
                "Description is required when uploading a new test.")

        test_dict['description'] = description
        if scival:
            test_dict['tags'] = ['scival']
        response = self.post(
            self._testapi,
            data=json.dumps(test_dict),
        )
        location = response.headers['location']
        return location.replace(common.BASE_URL, '')