fix: with args in flow.save_config (#3256)

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'] = []
 

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:

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