def test_memory_supervisor_organize():
supervisor = Default.get_employee('supervisor')
assert isinstance(supervisor.messenger, MemoryMessenger)
assert supervisor.messenger.topic == 'codepack'
assert supervisor.messenger.queues['codepack'].empty()
code1 = Code(add2)
code2 = Code(mul2)
code1 >> code2
code2(3, b=5)
assert code1.get_state() == 'UNKNOWN'
assert code2.get_state() == 'WAITING'
supervisor.organize()
assert code1.get_state() == 'UNKNOWN'
assert code2.get_state() == 'WAITING'
assert supervisor.messenger.queues['codepack'].empty()
code1(a=2, b=4)
assert code1.get_state() == 'TERMINATED'
assert code2.get_state() == 'WAITING'
assert supervisor.messenger.queues['codepack'].empty()
supervisor.organize()
assert code1.get_state() == 'TERMINATED'
assert code2.get_state() == 'READY'
assert not supervisor.messenger.queues['codepack'].empty(
) and supervisor.messenger.queues['codepack'].qsize() == 1
item = supervisor.messenger.queues['codepack'].get(block=False)
item.pop('source', None)
expected_item = code2.to_snapshot(args=(3, ),
kwargs={
'b': 5
},
timestamp=item['timestamp']).to_dict()
expected_item.pop('source', None)
assert item == expected_item
supervisor.close()
def test_get_default_worker_with_os_env(mock_kafka_consumer,
mock_docker_client):
try:
os.environ['CODEPACK_CONFIG_DIR'] = 'config'
os.environ['CODEPACK_CONFIG_PATH'] = 'sample.ini'
worker = Default.get_employee('worker')
assert isinstance(worker, Worker)
arg_list = mock_kafka_consumer.call_args_list
assert len(arg_list) == 1
args, kwargs = arg_list[0]
assert kwargs.get('bootstrap_servers', '') == '?:9092,?:9092,?:9092'
assert isinstance(kwargs.get('value_deserializer', ''), Callable)
assert len(args) == 1 and args[0] == 'codepack'
assert hasattr(worker.messenger, 'consumer')
assert getattr(worker.messenger,
'consumer').session is mock_kafka_consumer()
mock_docker_client.assert_called_once_with(
base_url='unix://var/run/docker.sock')
assert worker.docker_manager.docker.session == mock_docker_client()
assert isinstance(worker.callback_service, CallbackService)
assert isinstance(worker.logger, logging.Logger)
assert worker.logger.name == 'worker-logger'
assert hasattr(worker, 'script') and worker.script == 'run_snapshot.py'
default_dir = Default.get_config_instance().get_default_config_dir()
script_dir = os.path.join(default_dir, 'scripts')
assert hasattr(worker,
'script_dir') and worker.script_dir == script_dir
assert hasattr(worker,
'script_path') and worker.script_path == os.path.join(
default_dir, 'scripts/run_snapshot.py')
finally:
os.environ.pop('CODEPACK_CONFIG_DIR', None)
os.environ.pop('CODEPACK_CONFIG_PATH', None)
def test_memory_worker_run_snapshot_with_image_and_callback(
mock_docker_client):
worker = Default.get_employee('worker')
worker.callback = dummy_callback_function
assert isinstance(worker.messenger, MemoryMessenger)
code = Code(add2, image='dummy')
assert code.get_state() == 'UNKNOWN'
sn = worker.run_snapshot(code.to_snapshot(kwargs={'a': 3, 'b': 5}))
assert sn == code.serial_number
default_config_dir = Config.get_default_config_dir()
script_dir = os.path.join(default_config_dir, 'scripts')
mock_docker_client.return_value.containers.run.assert_called_once_with(
auto_remove=True,
command=[
'python', 'run_snapshot.py',
'%s.json' % sn, '-p', '.', '-l', 'worker-logger', '-c',
'dummy_callback_function'
],
dns=['8.8.8.8'],
environment=['CODEPACK_LOGGER_LOG_DIR=/usr/logs'],
image='dummy',
name=id(worker.docker_manager),
volumes=[
'%s:/usr/src/codepack' % script_dir,
'%s:/usr/logs' % os.path.abspath(Config.get_log_dir())
],
working_dir='/usr/src/codepack')
worker.stop()
def test_memory_worker_run_snapshot_with_nothing():
worker = Default.get_employee('worker')
assert isinstance(worker.messenger, MemoryMessenger)
code = Code(add2)
assert code.get_state() == 'UNKNOWN'
sn = worker.run_snapshot(code.to_snapshot(kwargs={'a': 3, 'b': 5}))
assert sn == code.serial_number
assert code.get_state() == 'TERMINATED'
assert code.get_result() == 8
worker.stop()
async def startup():
Default.init()
config = Default.config.get_config(section='apiserver')
_scheduler = config.get('scheduler', 'self')
common.add(key='supervisor', value=Default.get_employee('supervisor'), destroy=destroy_supervisor)
if isinstance(_scheduler, Scheduler):
common.add(key='scheduler', value=_scheduler, destroy=destroy_scheduler)
common.scheduler.start()
elif isinstance(_scheduler, str):
if _scheduler == 'self':
common.add(key='scheduler', value=Default.get_scheduler(), destroy=destroy_scheduler)
common.scheduler.start()
else:
common.add(key='scheduler', value=_scheduler)
def test_get_default_supervisor():
try:
os.environ['CODEPACK_CODESNAPSHOT_SOURCE'] = 'mongodb'
os.environ['CODEPACK_CODESNAPSHOT_DB'] = 'test_db'
os.environ['CODEPACK_CODESNAPSHOT_COLLECTION'] = 'test_collection'
supervisor = Default.get_employee('supervisor')
assert isinstance(supervisor, Supervisor)
assert isinstance(supervisor.messenger, MemoryMessenger)
assert supervisor.messenger.topic == 'codepack'
assert isinstance(supervisor.snapshot_service, SnapshotService)
assert isinstance(supervisor.snapshot_service.storage, MongoStorage)
finally:
os.environ.pop('CODEPACK_CODESNAPSHOT_SOURCE', None)
os.environ.pop('CODEPACK_CODESNAPSHOT_DB', None)
os.environ.pop('CODEPACK_CODESNAPSHOT_COLLECTION', None)
def test_memory_supervisor_run_code():
supervisor = Default.get_employee('supervisor')
assert isinstance(supervisor.messenger, MemoryMessenger)
assert supervisor.messenger.topic == 'codepack'
assert supervisor.messenger.queues['codepack'].empty()
code = Code(add2)
sn = supervisor.run_code(code=code, args=(3, ), kwargs={'b': 5})
assert sn == code.serial_number
assert not supervisor.messenger.queues['codepack'].empty(
) and supervisor.messenger.queues['codepack'].qsize() == 1
item = supervisor.messenger.queues['codepack'].get(block=False)
assert item == code.to_snapshot(args=(3, ),
kwargs={
'b': 5
},
timestamp=item['timestamp']).to_dict()
supervisor.close()
def test_memory_worker_run_snapshot_with_env(mock_subprocess_run):
worker = Default.get_employee('worker')
assert isinstance(worker.messenger, MemoryMessenger)
code = Code(add2, env='test_env', image='dummy')
assert code.get_state() == 'UNKNOWN'
sn = worker.run_snapshot(code.to_snapshot(kwargs={'a': 3, 'b': 5}))
assert sn == code.serial_number
default_config_dir = Config.get_default_config_dir()
script_dir = os.path.join(default_config_dir, 'scripts')
mock_subprocess_run.assert_called_once_with([
os.path.join(worker.interpreter_manager.path, 'test_env', 'bin',
'python'),
os.path.join(script_dir, 'run_snapshot.py'),
os.path.join(script_dir, '%s.json' % sn), '-p', script_dir, '-l',
'worker-logger'
])
worker.stop()
def test_get_default_worker_without_os_env(mock_docker_client):
worker = Default.get_employee('worker')
assert isinstance(worker, Worker)
assert isinstance(worker.messenger, MemoryMessenger)
assert worker.messenger.topic == 'codepack'
assert worker.supervisor is None
mock_docker_client.assert_called_once_with(
base_url='unix://var/run/docker.sock')
assert worker.docker_manager.docker.session == mock_docker_client()
assert isinstance(worker.callback_service, CallbackService)
assert isinstance(worker.callback_service.storage, FileStorage)
default_dir = Default.get_config_instance().get_default_config_dir()
assert worker.callback_service.storage.path == os.path.join(
default_dir, 'scripts')
assert isinstance(worker.logger, logging.Logger)
assert worker.logger.name == 'worker-logger'
assert hasattr(worker, 'script') and worker.script == 'run_snapshot.py'
assert hasattr(worker, 'script_dir') and worker.script_dir == os.path.join(
default_dir, 'scripts')
assert hasattr(worker,
'script_path') and worker.script_path == os.path.join(
default_dir, 'scripts/run_snapshot.py')
def test_memory_supervisor_run_codepack():
supervisor = Default.get_employee('supervisor')
assert isinstance(supervisor.messenger, MemoryMessenger)
assert supervisor.messenger.topic == 'codepack'
assert supervisor.messenger.queues['codepack'].empty()
code1 = Code(add2)
code2 = Code(mul2)
code3 = Code(combination)
code1 >> code3
code2 >> code3
code3.receive('c') << code1
code3.receive('d') << code2
codepack = CodePack('test', code=code1, subscribe=code3)
argpack = codepack.make_argpack()
argpack['add2'](a=3, b=5)
argpack['mul2'](a=2, b=3)
argpack['combination'](a=2, b=4)
sn = supervisor.run_codepack(codepack=codepack, argpack=argpack)
assert sn == codepack.serial_number
assert not supervisor.messenger.queues['codepack'].empty(
) and supervisor.messenger.queues['codepack'].qsize() == 3
items = [
supervisor.messenger.queues['codepack'].get(block=False)
for _ in range(3)
]
snapshots = [
c.to_snapshot(kwargs=argpack[i].to_dict()).to_dict()
for i, c in codepack.codes.items()
]
for item in items:
item.pop('timestamp', None)
for snapshot in snapshots:
snapshot.pop('timestamp', None)
assert sorted(items,
key=lambda x: x['id']) == sorted(snapshots,
key=lambda x: x['id'])
supervisor.close()
def test_get_default_supervisor_with_os_env(mock_kafka_producer):
try:
os.environ['CODEPACK_CONFIG_PATH'] = 'config/test.ini'
os.environ['CODEPACK_CODESNAPSHOT_SOURCE'] = 'mongodb'
os.environ['CODEPACK_CODESNAPSHOT_DB'] = 'test_db'
os.environ['CODEPACK_CODESNAPSHOT_COLLECTION'] = 'test_collection'
supervisor = Default.get_employee('supervisor')
assert isinstance(supervisor, Supervisor)
arg_list = mock_kafka_producer.call_args_list
assert len(arg_list) == 1
args, kwargs = arg_list[0]
assert kwargs.get('bootstrap_servers', '') == 'localhost:9092'
assert isinstance(kwargs.get('value_serializer', ''), Callable)
assert len(args) == 0
assert hasattr(supervisor.messenger, 'producer')
assert getattr(supervisor.messenger,
'producer').session is mock_kafka_producer()
assert isinstance(supervisor.snapshot_service, SnapshotService)
assert isinstance(supervisor.snapshot_service.storage, MongoStorage)
finally:
os.environ.pop('CODEPACK_CONFIG_PATH', None)
os.environ.pop('CODEPACK_CODESNAPSHOT_SOURCE', None)
os.environ.pop('CODEPACK_CODESNAPSHOT_DB', None)
os.environ.pop('CODEPACK_CODESNAPSHOT_COLLECTION', None)
from codepack import Default
if __name__ == '__main__':
worker = None
try:
worker = Default.get_employee('worker')
worker.start()
finally:
if worker:
worker.stop()
