def test_do_label_detection(make_stubber, stub_runner, make_labels,
monkeypatch, poll_status, error_code,
stop_on_method):
rekognition_client = boto3.client('rekognition')
rekognition_stubber = make_stubber(rekognition_client)
job_id = 'test-job-id'
job_status = 'TESTING'
labels = [
RekognitionLabel(label, time.time_ns()) for label in make_labels(3)
]
video = mock_video(monkeypatch, poll_status, rekognition_client)
with stub_runner(error_code, stop_on_method) as runner:
runner.add(rekognition_stubber.stub_start_detection,
'start_label_detection', video.video,
video.get_notification_channel(), job_id)
if poll_status == 'SUCCEEDED':
runner.add(rekognition_stubber.stub_get_label_detection, job_id,
job_status, labels)
if error_code is None:
got_labels = video.do_label_detection()
if poll_status == 'SUCCEEDED':
assert ([label.to_dict() for label in labels
] == [label.to_dict() for label in got_labels])
else:
assert got_labels == []
else:
with pytest.raises(ClientError) as exc_info:
video.do_label_detection()
assert exc_info.value.response['Error']['Code'] == error_code
def detect_labels(self, max_labels):
try:
response = self.rekognition_client.detect_labels(
Image=self.image, MaxLabels=max_labels)
labels = [RekognitionLabel(label) for label in response['Labels']]
logger.info("Found %s labels in %s.", len(labels), self.image_name)
except ClientError:
logger.info("Couldn't detect labels in %s.", self.image_name)
raise
else:
return labels
def do_label_detection(self):
"""
Performs label detection on the video.
:return: The list of labels found in the video.
"""
return self._do_rekognition_job(
"label detection",
self.rekognition_client.start_label_detection,
self.rekognition_client.get_label_detection,
lambda response: [
RekognitionLabel(label['Label'], label['Timestamp']) for label in
response['Labels']])
def detect_labels(self, max_labels):
"""
Detects labels in the image. Labels are objects and people.
:param max_labels: The maximum number of labels to return.
:return: The list of labels detected in the image.
"""
try:
response = self.rekognition_client.detect_labels(
Image=self.image, MaxLabels=max_labels)
labels = [RekognitionLabel(label) for label in response['Labels']]
logger.info("Found %s labels in %s.", len(labels), self.image_name)
except ClientError:
logger.info("Couldn't detect labels in %s.", self.image_name)
raise
else:
return labels
def test_detect_labels(make_stubber, make_labels, error_code):
rekognition_client = boto3.client('rekognition')
rekognition_stubber = make_stubber(rekognition_client)
image = RekognitionImage(TEST_IMAGE, 'test-image', rekognition_client)
labels = [RekognitionLabel(label) for label in make_labels(3)]
max_labels = 3
rekognition_stubber.stub_detect_labels(image.image,
max_labels,
labels,
error_code=error_code)
if error_code is None:
got_labels = image.detect_labels(max_labels)
assert ([label.to_dict() for label in labels
] == [label.to_dict() for label in got_labels])
else:
with pytest.raises(ClientError) as exc_info:
image.detect_labels(max_labels)
assert exc_info.value.response['Error']['Code'] == error_code