def test_lease_compare_result_to_status():
# Test the implicit conversion between CompareResult enum and LeaseUseResult status enum.
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.SAME, False) == LeaseProto.LeaseUseResult.STATUS_OK
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.SUPER_LEASE,
False) == LeaseProto.LeaseUseResult.STATUS_OLDER
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.SUPER_LEASE,
True) == LeaseProto.LeaseUseResult.STATUS_OK
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.SUB_LEASE,
False) == LeaseProto.LeaseUseResult.STATUS_OK
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.NEWER,
False) == LeaseProto.LeaseUseResult.STATUS_OK
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.OLDER,
False) == LeaseProto.LeaseUseResult.STATUS_OLDER
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.DIFFERENT_RESOURCES,
False) == LeaseProto.LeaseUseResult.STATUS_UNMANAGED
assert Lease.compare_result_to_lease_use_result_status(
Lease.CompareResult.DIFFERENT_EPOCHS,
False) == LeaseProto.LeaseUseResult.STATUS_WRONG_EPOCH
with pytest.raises(Exception):
Lease.compare_result_to_lease_use_result_status(100, False)
def test_bad_lease_constructors():
with pytest.raises(ValueError) as e_info:
Lease(None)
with pytest.raises(ValueError) as e_info:
lease_proto = LeaseProto.Lease()
Lease(lease_proto)
with pytest.raises(ValueError) as e_info:
lease_proto = LeaseProto.Lease()
lease_proto.resource = LLAMA
Lease(lease_proto)
with pytest.raises(ValueError) as e_info:
lease_proto = LeaseProto.Lease()
lease_proto.sequence[:] = SEQ
Lease(lease_proto)
def _sublease_or_none(self, leases, response, error_code):
# Look for the lease that we want.
matching_leases = [
lease for lease in leases if lease.resource == self.RESOURCE
]
# There should be exactly one match.
if len(matching_leases) == 1:
provided_lease = Lease(matching_leases[0])
return provided_lease.create_sublease()
# If there are NO leases, that's definitely an error.
if not matching_leases:
response.status = error_code
response.missing_lease_resources.append(self.RESOURCE)
return None
# Provided more than one lease for the resource!
response.header.error.code = header_pb2.CommonError.CODE_INVALID_REQUEST
response.header.error.message = '{} leases on resource {}'.format(
len(matching_leases), self.RESOURCE)
return None
def _create_lease(resource, epoch, sequence):
lease_proto = LeaseProto.Lease()
lease_proto.resource = resource
lease_proto.epoch = epoch
lease_proto.sequence[:] = sequence
return Lease(lease_proto)