def test_request_absolute_endpoint(): endpoint = "/api/service/health" full_url = f"http://localhost:9100{endpoint}" responses.add(responses.GET, full_url, json={}) client = Client(UsernamePasswordAuth("username", "password")) # If client does not properly handle absolute paths, client.get() will # raise a ConnectionRefused exception. client.get(endpoint)
def health_check(client: Client) -> bool: """ Query the health check API and check if each service is healthy (returns True) Args: client: the tamr client Returns: True if all services are healthy, False if unhealthy """ try: response = client.get(endpoint="/api/service/health") healthy_status = all( [value["healthy"] for value in response.json().values()]) if healthy_status: LOGGER.info( f"Client is healthy: {dumps(response.json(), indent=2)}") else: LOGGER.error( f"Client is unhealthy: {dumps(response.json(), indent=2)}") return healthy_status except requests.exceptions.ConnectionError as e: LOGGER.error(f"Could not connect to {client.host}. Error: {e}") return False
def get_with_connection_retry(client: Client, api_endpoint: str, *, timeout_seconds: int = 600, sleep_seconds: int = 20) -> requests.Response: """Will handle exceptions when attempting to connect to the Tamr API. This is used to handle connection issues when Tamr restarts due to a restore. Args: client: A Tamr client object api_endpoint: Tamr API endpoint timeout_seconds: Amount of time before a timeout error is thrown. Default is 600 seconds sleep_seconds: Amount of time in between attempts to connect to Tamr. Returns: A response object from API request.""" started = now() while timeout_seconds is None or now() - started < timeout_seconds: try: response = client.get(api_endpoint) return response except ConnectionError as e: # If we got for example a connection refused exception, try again later LOGGER.warning(f"Caught exception in connect {e}") sleep(sleep_seconds) raise TimeoutError( f"Took longer than {timeout_seconds} seconds to connect to tamr.")
def get_backup_by_id(client: Client, backup_id: str) -> JsonDict: """Fetches the json object for a given backup ID. Args: client: A Tamr client object. backup_id: The relativeID corresponding to the desired backup. Returns: Json dict corresponding to the desired backup. Raises: ValueError: Raised if GET request to Tamr fails """ api_string = f"backups/{backup_id}" response = client.get(api_string) if not response.ok: message = ( f"Received non-200 response code '{response.status_code}' " f"with message '{response.json()['message']}': '{response.json()}'" ) LOGGER.error(message) raise ValueError(message) return response.json()
def list_backups(client: Client) -> Generator[JsonDict, None, None]: """Lists all backups available to Tamr client. Will list both succeeded and failed backups. Args: client: A client object Returns: A generator of json dict objects for the backups available to client.""" response = client.get("backups") for backup in response.json(): yield backup
def from_resource_id(tamr: Client, *, job_id: Union[int, str]) -> Operation: """Create an operation from a job id Args: tamr: A Tamr client job_id: A job ID Returns: A Tamr operation """ job_response = tamr.get(f"/api/versioned/v1/operations/{job_id}") return Operation.from_response(tamr, job_response)
def test_request_session_cookie(): endpoint = "http://localhost:9100/api/versioned/v1/test" responses.add(responses.GET, endpoint, json={}) session = requests.Session() cookie = requests.cookies.create_cookie( name="test_cookie", value="the-cookie-works" ) session.cookies.set_cookie(cookie) client = Client(UsernamePasswordAuth("username", "password"), session=session) assert client.session is session endpoint = "test" client.get(endpoint) assert len(responses.calls) == 1 req = responses.calls[0].request assert req.url.endswith("test") assert req.headers.get("Cookie") is not None assert "test_cookie=" in req.headers.get("Cookie")
def current(client: Client) -> str: """Gets the version of Tamr for provided client Args: client: Tamr client Returns: String representation of Tamr version """ url = "/api/versioned/service/version" response = client.get(url).successful() return json.loads(response.content)["version"]
def get_all(tamr: Client) -> List[Operation]: """ Get a list of all jobs or operations. Args: tamr: A Tamr client Returns: A list of Operation objects. """ response = tamr.get( "/api/versioned/v1/operations", headers={"Accept": "application/json"}, stream=True ).json() ops = [Operation.from_json(tamr, item) for item in response] return ops
def test_base_path_no_base_path(): tamr = Client(auth) full_url = "http://localhost:9100/api/versioned/v1/datasets/2" responses.add(responses.GET, full_url, status=400) tamr.get("datasets/2")
def test_base_path_default_slash(): standard_base_path = "/api/versioned/v1/" tamr = Client(auth, base_path=standard_base_path) full_url = "http://localhost:9100/api/versioned/v1/datasets/1" responses.add(responses.GET, full_url, status=200) tamr.get("datasets/1")
def test_base_path_no_leading_slash(): bad_base_path = "api/versioned/v1/" tamr = Client(auth, base_path=bad_base_path) full_url = "http://localhost:9100/api/versioned/v1/datasets/1" responses.add(responses.GET, full_url, status=200) tamr.get("datasets/1")
def test__collect_operation_calls(): # setup mock client mock_client = Client(None) # setup mock operations base_operation_json = { "id": "2", "type": "SPARK", "description": "Profiling [employees_tiny.csv] attributes.", "status": { "state": "SUCCEEDED", "startTime": "2020-07-16T17:57:54.458Z", "endTime": "2020-07-16T17:58:22.836Z", "message": "", }, "created": { "username": "******", "time": "2020-07-16T17:57:28.920Z", "version": "82" }, "lastModified": { "username": "******", "time": "2020-07-16T17:58:23.977Z", "version": "119", }, "relativeId": "operations/2", } operation_states = [ OperationState.SUCCEEDED, OperationState.PENDING, OperationState.CANCELED, OperationState.RUNNING, OperationState.FAILED, ] mocks = {} for state in operation_states: op_json = base_operation_json.copy() op_json["status"]["state"] = state.value mock_operation = Operation.from_json(mock_client, op_json) mock_response = Response() mock_response._content = json.dumps(op_json).encode("utf-8") mock_response.status_code = 200 mocks[state] = {"op": mock_operation, "response": mock_response} # test succeeded with many pending mock_client.get = MagicMock(side_effect=[ # response while pending mocks[OperationState.PENDING]["response"], # polling mocks[OperationState.PENDING]["response"], mocks[OperationState.PENDING]["response"], mocks[OperationState.PENDING]["response"], mocks[OperationState.PENDING]["response"], mocks[OperationState.PENDING]["response"], mocks[OperationState.RUNNING]["response"], # response while running mocks[OperationState.RUNNING]["response"], # response while waiting mocks[OperationState.SUCCEEDED]["response"], # response when complete mocks[OperationState.SUCCEEDED]["response"], ]) with patch("tamr_toolbox.utils.client._from_response", return_value=mock_client): result_success = utils.testing._collect_operation_calls( response=mocks[OperationState.PENDING]["response"], poll_interval_seconds=0) assert len(result_success) == 3 for resp in result_success: assert resp.json()["id"] == "2" assert result_success[0].json( )["status"]["state"] == OperationState.PENDING.value assert result_success[1].json( )["status"]["state"] == OperationState.RUNNING.value assert result_success[2].json( )["status"]["state"] == OperationState.SUCCEEDED.value # test failed quickly mock_client.get = MagicMock(side_effect=[ # response while pending mocks[OperationState.PENDING]["response"], # polling mocks[OperationState.FAILED]["response"], # response while running mocks[OperationState.FAILED]["response"], # response while waiting mocks[OperationState.FAILED]["response"], # response when complete mocks[OperationState.FAILED]["response"], ]) with patch("tamr_toolbox.utils.client._from_response", return_value=mock_client): result_failed = utils.testing._collect_operation_calls( response=mocks[OperationState.PENDING]["response"], poll_interval_seconds=0) assert len(result_failed) == 3 for resp in result_failed: assert resp.json()["id"] == "2" assert result_failed[0].json( )["status"]["state"] == OperationState.PENDING.value assert result_failed[1].json( )["status"]["state"] == OperationState.FAILED.value assert result_failed[2].json( )["status"]["state"] == OperationState.FAILED.value
def test_base_path_no_trailing_slash(): bad_base_path = "/api/versioned/v1" unify = Client(auth, base_path=bad_base_path) full_url = "http://localhost:9100/api/versioned/v1/datasets/1" responses.add(responses.GET, full_url, status=200) unify.get("datasets/1")