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