def test_port_configuration(replicas_and_shards): def validate_ports_replica(shard, replica_port_in, replica_port_out, shards): assert replica_port_in == shard.args.port_in assert shard.args.port_out == replica_port_out peas_args = shard.peas_args peas = peas_args['peas'] assert len(peas) == shards if shards == 1: assert peas_args['head'] is None assert peas_args['tail'] is None assert peas[0].port_in == replica_port_in assert peas[0].port_out == replica_port_out else: shard_head = peas_args['head'] shard_tail = peas_args['tail'] assert shard.args.port_in == shard_head.port_in assert shard.args.port_out == shard_tail.port_out for pea in peas: assert shard_head.port_out == pea.port_in assert pea.port_out == shard_tail.port_in flow = Flow() for i, (replicas, shards) in enumerate(replicas_and_shards): flow.add( name=f'pod{i}', replicas=replicas, shards=shards, copy_flow=False, ) with flow: pods = flow._pod_nodes for pod_name, pod in pods.items(): if pod_name == 'gateway': continue if pod.args.replicas == 1: if int(pod.args.shards) == 1: assert len(pod.peas_args['peas']) == 1 else: assert len(pod.peas_args) == 3 shard_port_in = pod.args.port_in shard_port_out = pod.args.port_out else: shard_port_in = pod.head_args.port_out shard_port_out = pod.tail_args.port_in assert pod.head_args.port_in == pod.args.port_in assert pod.head_args.port_out == shard_port_in assert pod.tail_args.port_in == shard_port_out assert pod.tail_args.port_out == pod.args.port_out if pod.args.shards > 1: for shard in pod.shards: validate_ports_replica( shard, shard_port_in, shard_port_out, getattr(pod.args, 'replicas', 1), ) assert pod
def test_flow_to_k8s_yaml_external_pod(tmpdir, has_external): flow = Flow(name='test-flow', port=8080).add(name='executor0', ) if has_external: flow = flow.add(name='external_executor', external=True, host='1.2.3.4', port=9090) else: flow = flow.add(name='external_executor') dump_path = os.path.join(str(tmpdir), 'test_flow_k8s') namespace = 'test-flow-ns' flow.to_kubernetes_yaml( output_base_path=dump_path, k8s_namespace=namespace, ) yaml_dicts_per_deployment = { 'gateway': [], 'executor0': [], } assert len(set(os.listdir(dump_path))) == 2 if has_external else 3 for pod_name in set(os.listdir(dump_path)): file_set = set(os.listdir(os.path.join(dump_path, pod_name))) for file in file_set: with open(os.path.join(dump_path, pod_name, file)) as f: yml_document_all = list(yaml.safe_load_all(f)) yaml_dicts_per_deployment[file[:-4]] = yml_document_all gateway_objects = yaml_dicts_per_deployment['gateway'] gateway_args = gateway_objects[2]['spec']['template']['spec'][ 'containers'][0]['args'] assert ( gateway_args[gateway_args.index('--graph-description') + 1] == '{"executor0": ["external_executor"], "start-gateway": ["executor0"], "external_executor": ["end-gateway"]}' ) if has_external: assert '--deployments-addresses' in gateway_args assert ( gateway_args[gateway_args.index('--deployments-addresses') + 1] == '{"executor0": ["grpc://executor0.test-flow-ns.svc:8080"], "external_executor": ["grpc://1.2.3.4:9090"]}' ) else: assert '--deployments-addresses' in gateway_args assert ( gateway_args[gateway_args.index('--deployments-addresses') + 1] == '{"executor0": ["grpc://executor0.test-flow-ns.svc:8080"], "external_executor": ["grpc://external-executor.test-flow-ns.svc:8080"]}' )
def test_flow_equalities(): f1 = Flow().add().add(needs='gateway').needs_all(name='joiner') f2 = Flow().add(name='pod0').add(name='pod1', needs='gateway').add(name='joiner', needs=['pod0', 'pod1']) assert f1 == f2 f2 = f2.add() assert f1 != f2
def test_port_configuration(replicas_and_parallel): def extract_pod_args(pod): if 'replicas' not in pod.args or int(pod.args.replicas) == 1: head_args = pod.peas_args['head'] tail_args = pod.peas_args['tail'] middle_args = pod.peas_args['peas'] else: head_args = pod.head_args tail_args = pod.tail_args middle_args = pod.replicas_args return pod, head_args, tail_args, middle_args def get_outer_ports(pod, head_args, tail_args, middle_args): if not 'replicas' in pod.args or int(pod.args.replicas) == 1: if not 'parallel' in pod.args or int(pod.args.parallel) == 1: assert tail_args is None assert head_args is None replica = middle_args[0] # there is only one return replica.port_in, replica.port_out else: return pod.head_args.port_in, pod.tail_args.port_out else: assert pod.args.replicas == len(middle_args) return pod.head_args.port_in, pod.tail_args.port_out def validate_ports_pods(pods): for i in range(len(pods) - 1): _, port_out = get_outer_ports(*extract_pod_args(pods[i])) port_in_next, _ = get_outer_ports(*extract_pod_args(pods[i + 1])) assert port_out == port_in_next def validate_ports_replica(replica, replica_port_in, replica_port_out, parallel): assert replica_port_in == replica.args.port_in assert replica.args.port_out == replica_port_out peas_args = replica.peas_args peas = peas_args['peas'] assert len(peas) == parallel if parallel == 1: assert peas_args['head'] is None assert peas_args['tail'] is None assert peas[0].port_in == replica_port_in assert peas[0].port_out == replica_port_out else: shard_head = peas_args['head'] shard_tail = peas_args['tail'] assert replica.args.port_in == shard_head.port_in assert replica.args.port_out == shard_tail.port_out for pea in peas: assert shard_head.port_out == pea.port_in assert pea.port_out == shard_tail.port_in flow = Flow() for i, (replicas, parallel) in enumerate(replicas_and_parallel): flow.add( name=f'pod{i}', replicas=replicas, parallel=parallel, port_in=f'51{i}00', # info: needs to be set in this test since the test is asserting pod args with pod tail args port_out= f'51{i + 1}00', # outside this test, it don't have to be set copy_flow=False, ) with flow: pods = flow._pod_nodes validate_ports_pods( [pods['gateway']] + [pods[f'pod{i}'] for i in range(len(replicas_and_parallel))] + [pods['gateway']]) for pod_name, pod in pods.items(): if pod_name == 'gateway': continue if pod.args.replicas == 1: if int(pod.args.parallel) == 1: assert len(pod.peas_args['peas']) == 1 else: assert len(pod.peas_args) == 3 replica_port_in = pod.args.port_in replica_port_out = pod.args.port_out else: replica_port_in = pod.head_args.port_out replica_port_out = pod.tail_args.port_in assert pod.head_pea.args.port_in == pod.args.port_in assert pod.head_pea.args.port_out == replica_port_in assert pod.tail_pea.args.port_in == replica_port_out assert pod.tail_pea.args.port_out == pod.args.port_out if pod.args.replicas > 1: for replica in pod.replicas: validate_ports_replica( replica, replica_port_in, replica_port_out, getattr(pod.args, 'parallel', 1), ) assert pod