def test_request_id_reuse(self):
lease_request_1 = rpc_to_json(
rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.WINDOWS, ),
duration=7,
request_id='qwerty',
))
lease_request_2 = rpc_to_json(
rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.WINDOWS, ),
duration=189,
request_id='qwerty',
))
auth_testing.mock_get_current_identity(self)
lease_response_1 = jsonish_dict_to_rpc(
self.call_api('lease', lease_request_1).json,
rpc_messages.LeaseResponse,
)
lease_response_2 = jsonish_dict_to_rpc(
self.call_api('lease', lease_request_2).json,
rpc_messages.LeaseResponse,
)
self.failIf(lease_response_1.error)
self.assertEqual(
lease_response_2.error,
rpc_messages.LeaseRequestError.REQUEST_ID_REUSE,
)
self.assertNotEqual(
lease_response_1.request_hash,
lease_response_2.request_hash,
)
def test_one_request_one_matching_machine_entry(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX, ),
duration=1,
request_id='fake-id',
)
models.LeaseRequest(
deduplication_checksum=models.LeaseRequest.
compute_deduplication_checksum(request, ),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(client_request_id='fake-id', ),
).put()
models.CatalogMachineEntry.create_and_put(
rpc_messages.Dimensions(
backend=rpc_messages.Backend.DUMMY,
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
models.CatalogMachineEntryStates.AVAILABLE,
)
self.app.get(
'/internal/cron/process-lease-requests',
headers={'X-AppEngine-Cron': 'true'},
)
def test_add_duplicate(self):
request_1 = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(),
))
request_2 = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(),
))
self.mock_get_current_backend()
response_1 = jsonish_dict_to_rpc(
self.call_api('add_machine', request_1).json,
rpc_messages.CatalogManipulationResponse,
)
response_2 = jsonish_dict_to_rpc(
self.call_api('add_machine', request_2).json,
rpc_messages.CatalogManipulationResponse,
)
self.failIf(response_1.error)
self.assertEqual(
response_2.error,
rpc_messages.CatalogManipulationRequestError.HOSTNAME_REUSE,
)
def test_no_machine_id(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='fake-id',
)
lease_key = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
released=True,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
),
).put()
machine_key = models.CatalogMachineEntry(
id='id',
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
).put()
handlers_cron.release_lease(lease_key)
self.assertFalse(lease_key.get().response.lease_expiration_ts)
self.assertFalse(lease_key.get().released)
self.assertFalse(machine_key.get().lease_expiration_ts)
def test_no_expiration_ts(self):
self.mock(utils, 'enqueue_task', lambda *args, **kwargs: True)
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='fake-id',
)
lease_request_key = models.LeaseRequest(
id='id',
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
hostname='fake-host',
),
).put()
machine_key = models.CatalogMachineEntry(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
lease_id=lease_request_key.id(),
policies=rpc_messages.Policies(
machine_service_account='fake-service-account',
),
).put()
handlers_cron.reclaim_machine(machine_key, utils.utcnow())
self.assertTrue(lease_request_key.get().response.hostname)
def test_delete_invalid(self):
request_1 = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host-1',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
))
request_2 = rpc_to_json(
rpc_messages.CatalogMachineDeletionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host-2',
os_family=rpc_messages.OSFamily.LINUX,
), ))
request_3 = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host-1',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
))
self.mock_get_current_backend()
response_1 = jsonish_dict_to_rpc(
self.call_api('add_machine', request_1).json,
rpc_messages.CatalogManipulationResponse,
)
response_2 = jsonish_dict_to_rpc(
self.call_api('delete_machine', request_2).json,
rpc_messages.CatalogManipulationResponse,
)
response_3 = jsonish_dict_to_rpc(
self.call_api('add_machine', request_3).json,
rpc_messages.CatalogManipulationResponse,
)
self.failIf(response_1.error)
self.assertEqual(
response_2.error,
rpc_messages.CatalogManipulationRequestError.ENTRY_NOT_FOUND,
)
self.assertEqual(
response_3.error,
rpc_messages.CatalogManipulationRequestError.HOSTNAME_REUSE,
)
def test_releases(self):
self.mock(handlers_cron, 'release_lease', lambda *args, **kwargs: True)
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='fake-id',
)
key = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
released=True,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
state=rpc_messages.LeaseRequestState.UNTRIAGED,
),
).put()
self.app.get(
'/internal/cron/process-lease-releases',
headers={'X-AppEngine-Cron': 'true'},
)
def test_add_batch(self):
request = rpc_to_json(
rpc_messages.CatalogMachineBatchAdditionRequest(requests=[
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host-1',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
),
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host-2',
os_family=rpc_messages.OSFamily.WINDOWS,
),
policies=rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
),
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host-1',
os_family=rpc_messages.OSFamily.OSX,
),
policies=rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
),
], ))
self.mock_get_current_backend()
response = jsonish_dict_to_rpc(
self.call_api('add_machines', request).json,
rpc_messages.CatalogBatchManipulationResponse,
)
self.assertEqual(len(response.responses), 3)
self.failIf(response.responses[0].error)
self.failIf(response.responses[1].error)
self.assertEqual(
response.responses[2].error,
rpc_messages.CatalogManipulationRequestError.HOSTNAME_REUSE,
)
def test_delete(self):
request_1 = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
))
request_2 = rpc_to_json(
rpc_messages.CatalogMachineDeletionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
), ))
request_3 = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.WINDOWS,
),
policies=rpc_messages.Policies(
backend_project='fake-project',
backend_topic='fake-topic',
),
))
self.mock_get_current_backend()
response_1 = jsonish_dict_to_rpc(
self.call_api('add_machine', request_1).json,
rpc_messages.CatalogManipulationResponse,
)
response_2 = jsonish_dict_to_rpc(
self.call_api('delete_machine', request_2).json,
rpc_messages.CatalogManipulationResponse,
)
response_3 = jsonish_dict_to_rpc(
self.call_api('add_machine', request_3).json,
rpc_messages.CatalogManipulationResponse,
)
self.failIf(response_1.error)
self.failIf(response_2.error)
self.failIf(response_3.error)
def test_exact_match(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
disk_gb=100,
num_cpus=2,
os_family=rpc_messages.OSFamily.LINUX,
),
)
entry = models.CatalogMachineEntry(
dimensions=rpc_messages.Dimensions(
disk_gb=100,
num_cpus=2,
os_family=rpc_messages.OSFamily.LINUX,
),
)
self.assertTrue(handlers_cron.can_fulfill(entry, request))
def test_leased_task_failed(self):
self.mock(utils, 'enqueue_task', lambda *args, **kwargs: False)
machine_key = models.CatalogMachineEntry(
id='machine-id',
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(
machine_service_account='service-account',
),
state=models.CatalogMachineEntryStates.AVAILABLE,
).put()
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=1,
request_id='request-id',
)
lease_request_key = models.LeaseRequest(
id='lease-id',
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='client-request-id',
state=rpc_messages.LeaseRequestState.UNTRIAGED,
),
).put()
with self.assertRaises(handlers_cron.TaskEnqueuingError):
handlers_cron.lease_machine(machine_key, lease_request_key.get())
self.assertFalse(lease_request_key.get().machine_id)
self.assertEqual(
lease_request_key.get().response.state,
rpc_messages.LeaseRequestState.UNTRIAGED,
)
self.assertFalse(machine_key.get().lease_id)
self.assertEqual(
machine_key.get().state,
models.CatalogMachineEntryStates.AVAILABLE,
)
def test_mismatch(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
disk_gb=100,
num_cpus=4,
os_family=rpc_messages.OSFamily.LINUX,
),
)
entry = models.CatalogMachineEntry(
dimensions=rpc_messages.Dimensions(
disk_gb=100,
hostname='fake-host',
memory_gb=8.0,
num_cpus=2,
os_family=rpc_messages.OSFamily.LINUX,
),
)
self.assertFalse(handlers_cron.can_fulfill(entry, request))
def test_lease_not_found(self):
machine_key = models.CatalogMachineEntry(
id='id',
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
).put()
handlers_cron.release_lease(ndb.Key(models.LeaseRequest, 'fake-request'))
self.assertFalse(machine_key.get().lease_expiration_ts)
def test_reclaim_immediately(self):
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
duration=0,
request_id='fake-id',
)
lease = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
),
)
dimensions = rpc_messages.Dimensions(
backend=rpc_messages.Backend.DUMMY,
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
)
machine = models.CatalogMachineEntry(
dimensions=dimensions,
key=models.CatalogMachineEntry.generate_key(dimensions),
lease_id=lease.key.id(),
lease_expiration_ts=datetime.datetime.utcfromtimestamp(1),
policies=rpc_messages.Policies(
machine_service_account='fake-service-account',
),
state=models.CatalogMachineEntryStates.AVAILABLE,
).put()
lease.machine_id = machine.id()
lease.put()
self.app.get(
'/internal/cron/process-machine-reclamations',
headers={'X-AppEngine-Cron': 'true'},
)
def test_one_request_one_matching_machine_entry_lease_expiration_ts(self):
ts = int(utils.time_time())
request = rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX,
),
lease_expiration_ts=ts,
request_id='fake-id',
)
key = models.LeaseRequest(
deduplication_checksum=
models.LeaseRequest.compute_deduplication_checksum(request),
key=models.LeaseRequest.generate_key(
auth_testing.DEFAULT_MOCKED_IDENTITY.to_bytes(),
request,
),
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=request,
response=rpc_messages.LeaseResponse(
client_request_id='fake-id',
state=rpc_messages.LeaseRequestState.UNTRIAGED,
),
).put()
dimensions = rpc_messages.Dimensions(
backend=rpc_messages.Backend.DUMMY,
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
)
models.CatalogMachineEntry(
key=models.CatalogMachineEntry.generate_key(dimensions),
dimensions=dimensions,
policies=rpc_messages.Policies(
machine_service_account='fake-service-account',
),
state=models.CatalogMachineEntryStates.AVAILABLE,
).put()
self.app.get(
'/internal/cron/process-lease-requests',
headers={'X-AppEngine-Cron': 'true'},
)
self.assertEqual(key.get().response.lease_expiration_ts, ts)
def test_modify(self):
request = rpc_to_json(
rpc_messages.CatalogCapacityModificationRequest(
count=1,
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.OSX, ),
))
self.mock_get_current_backend()
response = jsonish_dict_to_rpc(
self.call_api('modify_capacity', request).json,
rpc_messages.CatalogManipulationResponse,
)
self.failIf(response.error)
def test_reclaimed(self):
lease_key = models.LeaseRequest(
id='fake-id',
deduplication_checksum='checksum',
machine_id='fake-host',
owner=auth_testing.DEFAULT_MOCKED_IDENTITY,
request=rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(),
request_id='request-id',
),
response=rpc_messages.LeaseResponse(
client_request_id='request-id',
hostname='fake-host',
),
).put()
machine_key = models.CatalogMachineEntry(
dimensions=rpc_messages.Dimensions(),
lease_id=lease_key.id(),
).put()
handlers_queues.reclaim(machine_key)
self.assertFalse(lease_key.get().machine_id)
self.assertFalse(lease_key.get().response.hostname)
self.assertFalse(machine_key.get())
def test_lease(self):
lease_request = rpc_to_json(
rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX, ),
duration=1,
request_id='abc',
))
auth_testing.mock_get_current_identity(self)
lease_response = jsonish_dict_to_rpc(
self.call_api('lease', lease_request).json,
rpc_messages.LeaseResponse,
)
self.failIf(lease_response.error)
def test_get(self):
models.CatalogMachineEntry(
key=models.CatalogMachineEntry._generate_key('DUMMY', 'fake-host'),
dimensions=rpc_messages.Dimensions(hostname='fake-host'),
).put()
request = rpc_to_json(
rpc_messages.CatalogMachineRetrievalRequest(
hostname='fake-host', ))
self.mock_get_current_backend()
response = jsonish_dict_to_rpc(
self.call_api('get', request).json,
rpc_messages.CatalogMachineRetrievalResponse,
)
self.assertEqual(response.dimensions.hostname, 'fake-host')
def test_add(self):
request = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(),
))
self.mock_get_current_backend()
response = jsonish_dict_to_rpc(
self.call_api('add_machine', request).json,
rpc_messages.CatalogManipulationResponse,
)
self.failIf(response.error)
def test_add_no_hostname(self):
request = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.LINUX, ),
policies=rpc_messages.Policies(),
))
self.mock_get_current_backend()
response = jsonish_dict_to_rpc(
self.call_api('add_machine', request).json,
rpc_messages.CatalogManipulationResponse,
)
self.assertEqual(
response.error,
rpc_messages.CatalogManipulationRequestError.UNSPECIFIED_HOSTNAME,
)
def test_project_without_topic(self):
lease_request = rpc_to_json(
rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.WINDOWS, ),
duration=9,
pubsub_project='my-project',
request_id='123',
))
auth_testing.mock_get_current_identity(self)
lease_response = jsonish_dict_to_rpc(
self.call_api('lease', lease_request).json,
rpc_messages.LeaseResponse,
)
self.assertEqual(
lease_response.error,
rpc_messages.LeaseRequestError.UNSPECIFIED_TOPIC,
)
def test_invalid_topic(self):
lease_request = rpc_to_json(
rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.WINDOWS, ),
duration=9,
pubsub_topic='../../a-different-project/topics/my-topic',
request_id='123',
))
auth_testing.mock_get_current_identity(self)
lease_response = jsonish_dict_to_rpc(
self.call_api('lease', lease_request).json,
rpc_messages.LeaseResponse,
)
self.assertEqual(
lease_response.error,
rpc_messages.LeaseRequestError.INVALID_TOPIC,
)
def test_add_project_without_topic(self):
request = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(pubsub_project='my-project', ),
))
self.mock_get_current_backend()
response = jsonish_dict_to_rpc(
self.call_api('add_machine', request).json,
rpc_messages.CatalogManipulationResponse,
)
self.assertEqual(
response.error,
rpc_messages.CatalogManipulationRequestError.UNSPECIFIED_TOPIC,
)
def test_add_invalid_topic(self):
request = rpc_to_json(
rpc_messages.CatalogMachineAdditionRequest(
dimensions=rpc_messages.Dimensions(
hostname='fake-host',
os_family=rpc_messages.OSFamily.LINUX,
),
policies=rpc_messages.Policies(
pubsub_topic='../../a-different-project/topics/my-topic',
),
))
self.mock_get_current_backend()
response = jsonish_dict_to_rpc(
self.call_api('add_machine', request).json,
rpc_messages.CatalogManipulationResponse,
)
self.assertEqual(
response.error,
rpc_messages.CatalogManipulationRequestError.INVALID_TOPIC,
)
def test_duplicate(self):
lease_request = rpc_to_json(
rpc_messages.LeaseRequest(
dimensions=rpc_messages.Dimensions(
os_family=rpc_messages.OSFamily.OSX, ),
duration=3,
request_id='asdf',
))
auth_testing.mock_get_current_identity(self)
lease_response_1 = jsonish_dict_to_rpc(
self.call_api('lease', lease_request).json,
rpc_messages.LeaseResponse,
)
lease_response_2 = jsonish_dict_to_rpc(
self.call_api('lease', lease_request).json,
rpc_messages.LeaseResponse,
)
self.failIf(lease_response_1.error)
self.failIf(lease_response_2.error)
self.assertEqual(
lease_response_1.request_hash,
lease_response_2.request_hash,
)