diff --git a/jina/jaml/parsers/flow/v1.py b/jina/jaml/parsers/flow/v1.py index f8cf5bd024196..5f8cf701b78b7 100644 --- a/jina/jaml/parsers/flow/v1.py +++ b/jina/jaml/parsers/flow/v1.py @@ -86,6 +86,9 @@ def dump(self, data: 'Flow') -> Dict: if data._kwargs: r['with'] = data._kwargs + if data._common_kwargs: + r['with'] = dict(r.get('with', {}), **data._common_kwargs) + if data._pod_nodes: r['executors'] = [] diff --git a/tests/unit/flow/test_flow_yaml_parser.py b/tests/unit/flow/test_flow_yaml_parser.py index fb2bf76ab7414..eda8557cf7f59 100644 --- a/tests/unit/flow/test_flow_yaml_parser.py +++ b/tests/unit/flow/test_flow_yaml_parser.py @@ -8,6 +8,7 @@ from jina.excepts import BadYAMLVersion from jina import Flow from jina.jaml import JAML +from jina.enums import GatewayProtocolType from jina.jaml.parsers import get_supported_versions from jina.parsers.flow import set_flow_parser from jina.types.document.generators import from_ndarray @@ -66,6 +67,41 @@ def test_load_dump_load(tmpdir): f2.save_config(str(Path(tmpdir) / 'a1.yml')) +def test_load_modify_dump_load(tmpdir): + f: Flow = Flow.load_config('yaml/flow-gateway.yml') + # assert vars inside `with` + assert f._kwargs['name'] == 'abc' + assert f.port_expose == 12345 + assert f.protocol == GatewayProtocolType.HTTP + # assert executor args + assert f._pod_nodes['custom1'].args.uses == 'jinahub://CustomExecutor1' + assert f._pod_nodes['custom2'].args.uses == 'CustomExecutor2' + assert f._pod_nodes['custom2'].args.port_in == 23456 + assert f._pod_nodes['custom2'].args.port_out == 34567 + + # change args inside `with` + f.port_expose = 12346 + f.protocol = GatewayProtocolType.WEBSOCKET + # change executor args + f._pod_nodes['custom2'].args.port_in = 23457 + f._pod_nodes['custom2'].args.port_out = 34568 + + f.save_config(str(Path(tmpdir) / 'a0.yml')) + f1: Flow = Flow.load_config(str(Path(tmpdir) / 'a0.yml')) + + # assert args from original yaml + assert f1._kwargs['name'] == 'abc' + assert 'custom1' in f1._pod_nodes + assert 'custom2' in f1._pod_nodes + assert f1._pod_nodes['custom1'].args.uses == 'jinahub://CustomExecutor1' + assert f1._pod_nodes['custom2'].args.uses == 'CustomExecutor2' + # assert args modified in code + assert f1.port_expose == 12346 + assert f1.protocol == GatewayProtocolType.WEBSOCKET + assert f1._pod_nodes['custom2'].args.port_in == 23457 + assert f1._pod_nodes['custom2'].args.port_out == 34568 + + def test_load_flow_with_port(): f = Flow.load_config('yaml/test-flow-port.yml') with f: diff --git a/tests/unit/flow/yaml/flow-gateway.yml b/tests/unit/flow/yaml/flow-gateway.yml new file mode 100644 index 0000000000000..6a0df8fac9ac3 --- /dev/null +++ b/tests/unit/flow/yaml/flow-gateway.yml @@ -0,0 +1,14 @@ +jtype: Flow +version: 1 +with: + name: abc + port_expose: 12345 + protocol: http + abc: def +executors: + - name: custom1 + uses: jinahub://CustomExecutor1 + - name: custom2 + uses: CustomExecutor2 + port_in: 23456 + port_out: 34567