[DEV-2834] Call EAS APIs for hosting capacity calibration

changelog.md

@@ -5,6 +5,7 @@
 
 ### New Features
 * Update `ModelConfig` to contain an optional `transformer_tap_settings` field to specify a set of distribution transformer tap settings to be applied by the model-processor.
+* Added basic client method to run a hosting capacity calibration and method to query its status.
 
 ### Enhancements
 * Added work package config documentation.

src/zepben/eas/client/eas_client.py

@@ -528,3 +528,95 @@ async def async_upload_study(self, study: Study):
                 else:
                     response = await response.text()
                 return response
+
+    def run_hosting_capacity_calibration(self, calibration_name: str):
+        """
+        Send request to run hosting capacity calibration
+        :param calibration_name: A string representation of the calibration name
+        :return: The HTTP response received from the Evolve App Server after attempting to run the calibration
+        """
+        return get_event_loop().run_until_complete(self.async_run_hosting_capacity_calibration(calibration_name))
+
+    async def async_run_hosting_capacity_calibration(self, calibration_name: str):
+        """
+        Send asynchronous request to run hosting capacity calibration
+        :param calibration_name: A string representation of the calibration name
+        :return: The HTTP response received from the Evolve App Server after attempting to run the calibration
+        """
+        with warnings.catch_warnings():
+            if not self._verify_certificate:
+                warnings.filterwarnings("ignore", category=InsecureRequestWarning)
+            json = {
+                "query": """
+                    mutation runCalibration($calibrationName: String!) {
+                        runCalibration(calibrationName: $calibrationName)
+                    }
+                """,
+                "variables": {
+                    "calibrationName": calibration_name
+                }
+            }
+            if self._verify_certificate:
+                sslcontext = ssl.create_default_context(cafile=self._ca_filename)
+
+            async with self.session.post(
+                    construct_url(protocol=self._protocol, host=self._host, port=self._port, path="/api/graphql"),
+                    headers=self._get_request_headers(),
+                    json=json,
+                    ssl=sslcontext if self._verify_certificate else False
+            ) as response:
+                if response.ok:
+                    response = await response.json()
+                else:
+                    response = await response.text()
+                return response
+
+    def get_hosting_capacity_calibration_run(self, id: str):
+        """
+        Retrieve information of a hosting capacity calibration run
+        :param id: The calibration run ID
+        :return: The HTTP response received from the Evolve App Server after requesting calibration run info
+        """
+        return get_event_loop().run_until_complete(self.async_get_hosting_capacity_calibration_run(id))
+
+    async def async_get_hosting_capacity_calibration_run(self, id: str):
+        """
+        Retrieve information of a hosting capacity calibration run
+        :param id: The calibration run ID
+        :return: The HTTP response received from the Evolve App Server after requesting calibration run info
+        """
+        with warnings.catch_warnings():
+            if not self._verify_certificate:
+                warnings.filterwarnings("ignore", category=InsecureRequestWarning)
+            json = {
+                "query": """
+                    query getCalibrationRun($id: ID!) {
+                        getCalibrationRun(calibrationRunId: $id) {
+                            id
+                            name
+                            workflowId
+                            runId
+                            startAt
+                            completedAt
+                            status
+                        }
+                    }
+                """,
+                "variables": {
+                    "id": id
+                }
+            }
+            if self._verify_certificate:
+                sslcontext = ssl.create_default_context(cafile=self._ca_filename)
+
+            async with self.session.post(
+                    construct_url(protocol=self._protocol, host=self._host, port=self._port, path="/api/graphql"),
+                    headers=self._get_request_headers(),
+                    json=json,
+                    ssl=sslcontext if self._verify_certificate else False
+            ) as response:
+                if response.ok:
+                    response = await response.json()
+                else:
+                    response = await response.text()
+                return response

test/test_eas_client.py

@@ -3,6 +3,7 @@
 #  This Source Code Form is subject to the terms of the Mozilla Public
 #  License, v. 2.0. If a copy of the MPL was not distributed with this
 #  file, You can obtain one at https://mozilla.org/MPL/2.0/.
+import json
 import random
 import ssl
 import string
@@ -12,6 +13,7 @@
 import pytest
 import trustme
 from pytest_httpserver import HTTPServer
+from werkzeug import Response
 from zepben.auth import ZepbenTokenFetcher
 
 from zepben.eas import EasClient, Study
@@ -492,3 +494,103 @@ def test_raises_error_if_access_token_and_client_secret_configured(httpserver: H
         )
     assert "Incompatible arguments passed to connect to secured Evolve App Server. You cannot provide multiple types of authentication. When using an access_token, do not provide client_id, client_secret, username, password, or token_fetcher." in str(
         error_message_for_username.value)
+
+def hosting_capacity_run_calibration_request_handler(request):
+    actual_body = json.loads(request.data.decode())
+    query = " ".join(actual_body['query'].split())
+
+    assert  query == "mutation runCalibration($calibrationName: String!) { runCalibration(calibrationName: $calibrationName) }"
+    assert actual_body['variables'] == {"calibrationName": "TEST CALIBRATION"}
+
+    return Response(json.dumps({"result": "success"}), status=200, content_type="application/json")
+
+def test_run_hosting_capacity_calibration_no_verify_success(httpserver: HTTPServer):
+    eas_client = EasClient(
+        LOCALHOST,
+        httpserver.port,
+        verify_certificate=False
+    )
+
+    httpserver.expect_oneshot_request("/api/graphql").respond_with_handler(hosting_capacity_run_calibration_request_handler)
+    res = eas_client.run_hosting_capacity_calibration("TEST CALIBRATION")
+    httpserver.check_assertions()
+    assert res == {"result": "success"}
+
+
+def test_run_hosting_capacity_calibration_invalid_certificate_failure(ca: trustme.CA, httpserver: HTTPServer):
+    with trustme.Blob(b"invalid ca").tempfile() as ca_filename:
+        eas_client = EasClient(
+            LOCALHOST,
+            httpserver.port,
+            verify_certificate=True,
+            ca_filename=ca_filename
+        )
+
+        httpserver.expect_oneshot_request("/api/graphql").respond_with_json({"result": "success"})
+        with pytest.raises(ssl.SSLError):
+            eas_client.run_hosting_capacity_calibration("TEST CALIBRATION")
+
+
+def test_run_hosting_capacity_calibration_valid_certificate_success(ca: trustme.CA, httpserver: HTTPServer):
+    with ca.cert_pem.tempfile() as ca_filename:
+        eas_client = EasClient(
+            LOCALHOST,
+            httpserver.port,
+            verify_certificate=True,
+            ca_filename=ca_filename
+        )
+
+        httpserver.expect_oneshot_request("/api/graphql").respond_with_handler(hosting_capacity_run_calibration_request_handler)
+        res = eas_client.run_hosting_capacity_calibration("TEST CALIBRATION")
+        httpserver.check_assertions()
+        assert res == {"result": "success"}
+
+def get_hosting_capacity_run_calibration_request_handler(request):
+    actual_body = json.loads(request.data.decode())
+    query = " ".join(actual_body['query'].split())
+
+    assert  query == "query getCalibrationRun($id: ID!) { getCalibrationRun(calibrationRunId: $id) { id name workflowId runId startAt completedAt status } }"
+    assert actual_body['variables'] == {"id": "calibration-id"}
+
+    return Response(json.dumps({"result": "success"}), status=200, content_type="application/json")
+
+def test_get_hosting_capacity_calibration_run_no_verify_success(httpserver: HTTPServer):
+    eas_client = EasClient(
+        LOCALHOST,
+        httpserver.port,
+        verify_certificate=False
+    )
+
+    httpserver.expect_oneshot_request("/api/graphql").respond_with_handler(get_hosting_capacity_run_calibration_request_handler)
+    res = eas_client.get_hosting_capacity_calibration_run("calibration-id")
+    httpserver.check_assertions()
+    assert res == {"result": "success"}
+
+
+def test_get_hosting_capacity_calibration_run_invalid_certificate_failure(ca: trustme.CA, httpserver: HTTPServer):
+    with trustme.Blob(b"invalid ca").tempfile() as ca_filename:
+        eas_client = EasClient(
+            LOCALHOST,
+            httpserver.port,
+            verify_certificate=True,
+            ca_filename=ca_filename
+        )
+
+        httpserver.expect_oneshot_request("/api/graphql").respond_with_json({"result": "success"})
+        with pytest.raises(ssl.SSLError):
+            eas_client.get_hosting_capacity_calibration_run("calibration-id")
+
+
+def test_get_hosting_capacity_calibration_run_valid_certificate_success(ca: trustme.CA, httpserver: HTTPServer):
+    with ca.cert_pem.tempfile() as ca_filename:
+        eas_client = EasClient(
+            LOCALHOST,
+            httpserver.port,
+            verify_certificate=True,
+            ca_filename=ca_filename
+        )
+
+        httpserver.expect_oneshot_request("/api/graphql").respond_with_handler(get_hosting_capacity_run_calibration_request_handler)
+        res = eas_client.get_hosting_capacity_calibration_run("calibration-id")
+        httpserver.check_assertions()
+        assert res == {"result": "success"}