def test_TransportPublish(self):
port = self._freeTCPPort()
popen = subprocess.Popen(
["python", "test/twistedserver_publishperiodically.py",
str(port)])
self._waitForServerToBeReady(port)
global testEvent
testEvent = threading.Event()
def setEvent(arguments):
if arguments != dict(data="fake data"):
raise Exception("Bad arguments: %s" % (arguments, ))
global testEvent
testEvent.set()
try:
tested = subscribe.Subscribe("tcp://localhost:%d" % port)
try:
tested.register(setEvent)
testEvent.wait(1)
self.assertTrue(testEvent.isSet())
finally:
tested.close()
finally:
popen.terminate()
def _getAllocationInstance(self, requirements, allocationID):
if self._subscribe is None:
self._subscribe = subscribe.Subscribe(
amqpURL=self._providerSubscribeLocation,
skipSuicide=self._skipSuicide)
allocationInstance = allocation.Allocation(id=allocationID,
requirements=requirements,
ipcClient=self,
subscribe=self._subscribe,
heartbeat=self._heartbeat)
self._activeAllocations.append(allocationInstance)
return allocationInstance
def __init__(self, providerRequestLocation, providerSubscribeLocation,
providerHTTPLocation):
self._providerHTTPLocation = providerHTTPLocation
self._request = transport.Transport(providerRequestLocation)
self._lock = threading.Lock()
self._closed = False
self._activeAllocations = []
self.call("handshake",
versionInfo=dict(RACKATTACK_VERSION=api.VERSION))
self._subscribe = subscribe.Subscribe(
connectTo=providerSubscribeLocation)
self._connectionToProviderInterrupted = suicide.killSelf
self._heartbeat = heartbeat.HeartBeat(self)
def allocate(self, requirements, allocationInfo):
assert len(requirements) > 0
jsonableRequirements = {
name: requirement.__dict__
for name, requirement in requirements.iteritems()
}
if self._subscribe is None:
self._subscribe = subscribe.Subscribe(
amqpURL=self._providerSubscribeLocation,
skipSuicide=self._skipSuicide)
allocationID = self.call(cmd='allocate',
requirements=jsonableRequirements,
allocationInfo=allocationInfo.__dict__)
return self._getAllocationInstance(requirements, allocationID)
def main():
args = get_args()
config.configure_logger()
db = ElasticsearchDBMock()
if args.script == "allocations":
monitor = mock.Mock()
_, amqp_url, _ = os.environ['RACKATTACK_PROVIDER'].split("@@")
subscription_mgr = subscribe.Subscribe(amqp_url)
allocation_handler = AllocationsHandler(subscription_mgr, db, monitor)
allocation_handler.run()
elif args.script == "smart":
smart_scanner = smartscanner.SmartScanner(db)
smart_scanner.run()
else:
raise ValueError(args.script)
def _generate_instance_with_mocked_event_loop(self):
self._db = ElasticsearchDBMock()
subscription_mgr = subscribe.Subscribe(mock_pika.DEFAULT_AMQP_URL)
instance = AllocationsHandler(subscription_mgr, self._db, mock.Mock())
def queueGetWrapper(*args, **kwargs):
if self.tested._tasks.qsize() > 0:
item = self.original_get()
else:
item = greenlet.getcurrent().parent.switch()
return item
self.original_get = instance._tasks.get
instance._tasks.get = queueGetWrapper
return instance
def create_subscription():
_, amqp_url, _ = os.environ['RACKATTACK_PROVIDER'].split("@@")
subscription_mgr = subscribe.Subscribe(amqp_url)
return subscription_mgr