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")
