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