Add refresh default to canvas.update_process_group

Add convenience methods to parameters.py to assign and remove a Parameter Context to a Process Group, with appropriate tests
Minor coding style improvements

nipyapi/canvas.py

@@ -719,20 +719,25 @@ def _starting_schedule_processor(processor_):
         return False
 
 
-def update_process_group(pg, update):
+def update_process_group(pg, update, refresh=True):
     """
         Updates a given Process Group.
 
         Args:
-            pg (ProcessGroupEntity): The Processor to target for update
+            pg (ProcessGroupEntity): The Process Group to
+              target for update
             update (dict): key:value pairs to update
+            refresh (bool): Whether to refresh the Process Group before
+              applying the update
 
         Returns:
             (ProcessGroupEntity): The updated ProcessorEntity
 
         """
     assert isinstance(pg, nipyapi.nifi.ProcessGroupEntity)
     with nipyapi.utils.rest_exceptions():
+        if refresh:
+            pg = get_process_group(pg.id, 'id')
         return nipyapi.nifi.ProcessGroupsApi().update_process_group(
             id=pg.id,
             body=nipyapi.nifi.ProcessGroupEntity(
@@ -1282,6 +1287,7 @@ def get_controller(identifier, identifier_type='name', bool_response=False):
     assert isinstance(identifier, six.string_types)
     assert identifier_type in ['name', 'id']
     handle = nipyapi.nifi.ControllerServicesApi()
+    out = None
     try:
         if identifier_type == 'id':
             out = handle.get_controller_service(identifier)

nipyapi/config.py

@@ -25,7 +25,7 @@
 # convenience function for this in nipyapi.utils.set_endpoint
 
 # Set Default Host for NiFi
-default_host = '34.244.104.232'  # Default to localhost for release
+default_host = 'localhost'  # Default to localhost for release
 #
 nifi_config.host = os.getenv(
     'NIFI_API_ENDPOINT',

nipyapi/parameters.py

@@ -18,7 +18,9 @@
     "list_all_parameter_contexts", "create_parameter_context",
     "delete_parameter_context", "get_parameter_context",
     "update_parameter_context", "prepare_parameter",
-    "delete_parameter_from_context"
+    "delete_parameter_from_context", "upsert_parameter_to_context",
+    "assign_context_to_process_group",
+    "remove_context_from_process_group"
 ]
 
 
@@ -217,3 +219,45 @@ def upsert_parameter_to_context(context, parameter):
     enforce_min_ver('1.10.0')
     context.component.parameters = [parameter]
     return update_parameter_context(context=context)
+
+
+def assign_context_to_process_group(pg, context_id):
+    """
+    Assigns a given Parameter Context to a specific Process Group
+
+    Args:
+        pg (ProcessGroupEntity): The Process Group to target
+        context_id (str): The ID of the Parameter Context
+
+    Returns:
+        (ProcessGroupEntity) The updated Process Group
+    """
+    assert isinstance(context_id, str)
+    return nipyapi.canvas.update_process_group(
+        pg=pg,
+        update={
+            'parameter_context': {
+                'id': context_id
+            }
+        }
+    )
+
+
+def remove_context_from_process_group(pg):
+    """
+    Clears any Parameter Context from the given Process Group
+
+    Args:
+        pg (ProcessGroupEntity): The Process Group to target
+
+    Returns:
+        (ProcessGroupEntity) The updated Process Group
+    """
+    return nipyapi.canvas.update_process_group(
+        pg=pg,
+        update={
+            'parameter_context': {
+                'id': None
+            }
+        }
+    )

nipyapi/utils.py

@@ -573,7 +573,6 @@ def wrapper(*args, **kwargs):
                     nipyapi.registry.rest.ApiException) as e:
                 if status_code is not None and e.status == int(status_code):
                     return response
-                else:
-                    _raise(ValueError(e.body), e)
+                _raise(ValueError(e.body), e)
         return wrapper
     return func_wrapper

tests/test_parameters.py

@@ -111,3 +111,23 @@ def test_upsert_parameter_to_context(regress_nifi, fix_context):
         parameters.prepare_parameter('Black', 'Lives', 'Matter', True)
     )
     assert r1.component.parameters[0].parameter.description == 'Matter'
+
+
+def test_assign_context_to_process_group(regress_nifi, fix_pg, fix_context):
+    if check_version('1.10.0') > 0:
+        pytest.skip("NiFi not 1.10+")
+    c1 = fix_context.generate()
+    pg1 = fix_pg.generate()
+    r1 = parameters.assign_context_to_process_group(pg1, c1.id)
+    assert r1.component.parameter_context.id == c1.id
+
+
+def test_remove_context_from_process_group(regress_nifi, fix_pg, fix_context):
+    if check_version('1.10.0') > 0:
+        pytest.skip("NiFi not 1.10+")
+    c1 = fix_context.generate()
+    pg1 = fix_pg.generate()
+    r1 = parameters.assign_context_to_process_group(pg1, c1.id)
+    assert r1.component.parameter_context.id == c1.id
+    r2 = parameters.remove_context_from_process_group(pg1)
+    assert r2.component.parameter_context is None