chore: ruff formatting

Author: khpeet

newrelic/hooks/external_pyzeebe.py

@@ -26,54 +26,56 @@
 
 CLIENT_ATTRIBUTES_WARNING_LOG_MSG = "Exception occurred in PyZeebe instrumentation: Failed to add client attributes."
 
+
 # Adds client method params as txn or span attributes
 def _add_client_input_attributes(method_name, txn, args, kwargs):
     try:
         if method_name in ("run_process", "run_process_with_result"):
             bpmn_id = kwargs.get("bpmn_process_id", args[0] if args else None)
             if bpmn_id:
                 txn._add_agent_attribute("zeebe.client.bpmnProcessId", bpmn_id)
-                #add_attr("zeebe.client.bpmnProcessId", bpmn_id)
+                # add_attr("zeebe.client.bpmnProcessId", bpmn_id)
         elif method_name == "publish_message":
             msg_name = kwargs.get("name", args[0] if args else None)
             if msg_name:
                 txn._add_agent_attribute("zeebe.client.messageName", msg_name)
-                #add_attr("zeebe.client.messageName", msg_name)
+                # add_attr("zeebe.client.messageName", msg_name)
             correlation_key = kwargs.get("correlation_key", args[1] if args and len(args) > 1 else None)
             if correlation_key:
                 txn._add_agent_attribute("zeebe.client.correlationKey", correlation_key)
-                #add_attr("zeebe.client.correlationKey", correlation_key)
+                # add_attr("zeebe.client.correlationKey", correlation_key)
             message_id = kwargs.get("message_id")
             if message_id and len(args) > 4:
                 message_id = args[4]
             if message_id:
                 txn._add_agent_attribute("zeebe.client.messageId", message_id)
-                #add_attr("zeebe.client.messageId", message_id)
+                # add_attr("zeebe.client.messageId", message_id)
         elif method_name == "deploy_resource":
             resources = list(args)
             if len(resources) == 1 and isinstance(resources[0], (list, tuple)):
                 resources = list(resources[0])
             if resources:
                 txn._add_agent_attribute("zeebe.client.resourceCount", len(resources))
-                #add_attr("zeebe.client.resourceCount", len(resources))
+                # add_attr("zeebe.client.resourceCount", len(resources))
                 if len(resources) == 1:
                     try:
                         txn._add_agent_attribute("zeebe.client.resourceFile", str(resources[0]))
-                        #add_attr("zeebe.client.resourceFile", str(resources[0]))
+                        # add_attr("zeebe.client.resourceFile", str(resources[0]))
                     except Exception:
                         txn._add_agent_attribute("zeebe.client.resourceFile", str(resources[0]))
-                        #add_attr("zeebe.client.resourceFile", repr(resources[0]))
+                        # add_attr("zeebe.client.resourceFile", repr(resources[0]))
     except Exception:
         _logger.warning(CLIENT_ATTRIBUTES_WARNING_LOG_MSG, exc_info=True)
 
+
 # Async wrapper that instruments router/worker annotations`
 async def _nr_wrapper_execute_one_job(wrapped, instance, args, kwargs):
     job = args[0] if args else kwargs.get("job")
     process_id = getattr(job, "bpmn_process_id", None) or "UnknownProcess"
     task_type = getattr(job, "type", None) or "UnknownType"
     txn_name = f"{process_id}/{task_type}"
 
-    with BackgroundTask(application_instance(), txn_name, group="ZeebeTask"): 
+    with BackgroundTask(application_instance(), txn_name, group="ZeebeTask"):
         if job is not None:
             if hasattr(job, "key"):
                 add_custom_attribute("zeebe.job.key", job.key)
@@ -89,6 +91,7 @@ async def _nr_wrapper_execute_one_job(wrapped, instance, args, kwargs):
         result = await wrapped(*args, **kwargs)
         return result
 
+
 # Async wrapper that instruments a ZeebeClient method.
 def _nr_client_wrapper(method_name):
     async def wrapper(wrapped, instance, args, kwargs):
@@ -108,34 +111,23 @@ async def wrapper(wrapped, instance, args, kwargs):
                 result = await wrapped(*args, **kwargs)
             finally:
                 created_txn.__exit__(None, None, None)
-            
+
             return result
 
     return wrapper
 
+
 # Instrument JobExecutor.execute_one_job to create a background transaction per job (invoked from @router.task or @worker.task annotations)
 def instrument_pyzeebe_worker_job_executor(module):
     if hasattr(module, "JobExecutor"):
-        wrap_function_wrapper(
-            module,
-            "JobExecutor.execute_one_job",
-            _nr_wrapper_execute_one_job
-        )
+        wrap_function_wrapper(module, "JobExecutor.execute_one_job", _nr_wrapper_execute_one_job)
+
 
 # Instrument ZeebeClient methods to trace client calls.
 def instrument_pyzeebe_client_client(module):
-    target_methods = (
-        "run_process",
-        "run_process_with_result",
-        "deploy_resource",
-        "publish_message",
-    )
+    target_methods = ("run_process", "run_process_with_result", "deploy_resource", "publish_message")
 
     for method_name in target_methods:
         if hasattr(module, "ZeebeClient"):
             if hasattr(module.ZeebeClient, method_name):
-                wrap_function_wrapper(
-                    module,
-                    f"ZeebeClient.{method_name}",
-                    _nr_client_wrapper(method_name)
-                )
+                wrap_function_wrapper(module, f"ZeebeClient.{method_name}", _nr_client_wrapper(method_name))

tests/external_pyzeebe/_mocks.py

@@ -2,74 +2,86 @@
 from pyzeebe.grpc_internals.zeebe_adapter import ZeebeAdapter
 
 # Dummy response objects with only required fields
-DummyCreateProcessInstanceResponse = SimpleNamespace(
-    process_instance_key=12345
-)
+DummyCreateProcessInstanceResponse = SimpleNamespace(process_instance_key=12345)
 
 DummyCreateProcessInstanceWithResultResponse = SimpleNamespace(
-    process_instance_key=45678,
-    variables={"result": "success"}
+    process_instance_key=45678, variables={"result": "success"}
 )
 
-DummyDeployResourceResponse = SimpleNamespace(
-    key=67890,
-    deployments=[],
-    tenant_id=None
-)
+DummyDeployResourceResponse = SimpleNamespace(key=67890, deployments=[], tenant_id=None)
+
+DummyPublishMessageResponse = SimpleNamespace(key=99999, tenant_id=None)
 
-DummyPublishMessageResponse = SimpleNamespace(
-    key=99999,
-    tenant_id=None
-)
 
 # Dummy RPC stub coroutines
-async def dummy_create_process_instance(self, bpmn_process_id: str, variables: dict = None, version: int = -1, tenant_id: str = None):
+async def dummy_create_process_instance(
+    self, bpmn_process_id: str, variables: dict = None, version: int = -1, tenant_id: str = None
+):
     """Simulate ZeebeAdapter.create_process_instance"""
     return DummyCreateProcessInstanceResponse
 
-async def dummy_create_process_instance_with_result(self, bpmn_process_id: str, variables: dict = None, version: int = -1, timeout: int = 0, variables_to_fetch=None, tenant_id: str = None):
+
+async def dummy_create_process_instance_with_result(
+    self,
+    bpmn_process_id: str,
+    variables: dict = None,
+    version: int = -1,
+    timeout: int = 0,
+    variables_to_fetch=None,
+    tenant_id: str = None,
+):
     """Simulate ZeebeAdapter.create_process_instance_with_result"""
     return DummyCreateProcessInstanceWithResultResponse
 
+
 async def dummy_deploy_resource(*resource_file_path: str, tenant_id: str = None):
     """Simulate ZeebeAdapter.deploy_resource"""
     # Create dummy deployment metadata for each provided resource path
     deployments = []
     for path in resource_file_path:
-        deployments.append(SimpleNamespace(
-            resource_name=str(path),
-            bpmn_process_id="dummy_process",
-            process_definition_key=123,
-            version=1,
-            tenant_id=tenant_id if tenant_id is not None else None
-        ))
+        deployments.append(
+            SimpleNamespace(
+                resource_name=str(path),
+                bpmn_process_id="dummy_process",
+                process_definition_key=123,
+                version=1,
+                tenant_id=tenant_id if tenant_id is not None else None,
+            )
+        )
     # Create a dummy response with a list of deployments
     return SimpleNamespace(
-        deployment_key=333333,
-        deployments=deployments,
-        tenant_id=tenant_id if tenant_id is not None else None
+        deployment_key=333333, deployments=deployments, tenant_id=tenant_id if tenant_id is not None else None
     )
 
-async def dummy_publish_message(self, name: str, correlation_key: str, variables: dict = None, time_to_live_in_milliseconds: int = 60000, message_id: str = None, tenant_id: str = None):
+
+async def dummy_publish_message(
+    self,
+    name: str,
+    correlation_key: str,
+    variables: dict = None,
+    time_to_live_in_milliseconds: int = 60000,
+    message_id: str = None,
+    tenant_id: str = None,
+):
     """Simulate ZeebeAdapter.publish_message"""
     # Return the dummy response (contains message key)
-    return SimpleNamespace(
-        key=999999,
-        tenant_id=tenant_id if tenant_id is not None else None
-    )
+    return SimpleNamespace(key=999999, tenant_id=tenant_id if tenant_id is not None else None)
+
 
 async def dummy_complete_job(self, job_key: int, variables: dict):
     """Simulate JobExecutor.complete_job"""
     setattr(self, "_last_complete", {"job_key": job_key, "variables": variables})
     return None
 
+
 class DummyZeebeAdapter(ZeebeAdapter):
     """Simulate a ZeebeAdapter so JobExecutor can be instatiated w/o gRPC channel"""
+
     def __init__(self):
         self.completed_job_key = None
         self.completed_job_vars = None
+
     async def complete_job(self, job_key: int, variables: dict):
         self.completed_job_key = job_key
         self.completed_job_vars = variables
         return None
-

tests/external_pyzeebe/conftest.py

@@ -28,4 +28,4 @@
 
 collector_agent_registration = collector_agent_registration_fixture(
     app_name="Python Agent Test (external_pyzeebe)", default_settings=_default_settings
-)
+)

tests/external_pyzeebe/test_client.py

@@ -16,7 +16,12 @@
 import asyncio
 from pyzeebe import ZeebeClient, create_insecure_channel
 from pyzeebe.grpc_internals.zeebe_adapter import ZeebeAdapter
-from _mocks import dummy_create_process_instance, dummy_create_process_instance_with_result, dummy_deploy_resource, dummy_publish_message
+from _mocks import (
+    dummy_create_process_instance,
+    dummy_create_process_instance_with_result,
+    dummy_deploy_resource,
+    dummy_publish_message,
+)
 
 from newrelic.api.background_task import background_task
 from testing_support.validators.validate_transaction_metrics import validate_transaction_metrics
@@ -25,12 +30,9 @@
 
 client = ZeebeClient(create_insecure_channel())
 
+
 @validate_transaction_metrics(
-    "test_zeebe_client:run_process",
-    rollup_metrics=[
-        ("ZeebeClient/run_process", 1)
-    ],
-    background_task=True
+    "test_zeebe_client:run_process", rollup_metrics=[("ZeebeClient/run_process", 1)], background_task=True
 )
 @validate_attributes("agent", ["zeebe.client.bpmnProcessId"])
 def test_run_process(monkeypatch):
@@ -40,15 +42,14 @@ def test_run_process(monkeypatch):
     def _test():
         response = asyncio.run(client.run_process("test_process"))
         assert response.process_instance_key == 12345
+
     _test()
 
 
 @validate_transaction_metrics(
     "test_zeebe_client:run_process_with_result",
-    rollup_metrics=[
-        ("ZeebeClient/run_process_with_result", 1)
-    ],
-    background_task=True
+    rollup_metrics=[("ZeebeClient/run_process_with_result", 1)],
+    background_task=True,
 )
 @validate_attributes("agent", ["zeebe.client.bpmnProcessId"])
 def test_run_process_with_result(monkeypatch):
@@ -59,15 +60,12 @@ def _test():
         result = asyncio.run(client.run_process_with_result("test_process"))
         assert result.process_instance_key == 45678
         assert result.variables == {"result": "success"}
+
     _test()
 
 
 @validate_transaction_metrics(
-    "test_zeebe_client:deploy_resource",
-    rollup_metrics=[
-        ("ZeebeClient/deploy_resource", 1)
-    ],
-    background_task=True
+    "test_zeebe_client:deploy_resource", rollup_metrics=[("ZeebeClient/deploy_resource", 1)], background_task=True
 )
 @validate_attributes("agent", ["zeebe.client.resourceCount", "zeebe.client.resourceFile"])
 def test_deploy_resource(monkeypatch):
@@ -77,15 +75,12 @@ def test_deploy_resource(monkeypatch):
     def _test():
         result = asyncio.run(client.deploy_resource("test.bpmn"))
         assert result.deployment_key == 333333
+
     _test()
 
 
 @validate_transaction_metrics(
-    "test_zeebe_client:publish_message",
-    rollup_metrics=[
-        ("ZeebeClient/publish_message", 1)
-    ],
-    background_task=True
+    "test_zeebe_client:publish_message", rollup_metrics=[("ZeebeClient/publish_message", 1)], background_task=True
 )
 @validate_attributes("agent", ["zeebe.client.messageName", "zeebe.client.correlationKey", "zeebe.client.messageId"])
 def test_publish_message(monkeypatch):
@@ -95,4 +90,5 @@ def test_publish_message(monkeypatch):
     def _test():
         result = asyncio.run(client.publish_message(name="test_message", correlation_key="999999", message_id="abc123"))
         assert result.key == 999999
-    _test()
+
+    _test()

tests/external_pyzeebe/test_job_executor.py

@@ -12,13 +12,14 @@
 # Set up a router with a dummy async task
 router = ZeebeTaskRouter()
 
+
 @router.task(task_type="testTask")
 async def dummy_task(x: int) -> dict:
     """
-    Simulate a task function that reads input variable 'x'
-    and adds 1.
+    Simulate a task function that reads input variable x
     """
-    return {"result": x }
+    return {"result": x}
+
 
 # @validate_transaction_metrics(
 #     "ZeebeTask/test_process/testTask",
@@ -27,20 +28,22 @@ async def dummy_task(x: int) -> dict:
 #     ],
 #     background_task=True
 # )
-@validate_custom_parameters(required_params=[
-    ("zeebe.job.key", 123),
-    ("zeebe.job.type", "testTask"),
-    ("zeebe.job.bpmnProcessId", "test_process"),
-    ("zeebe.job.processInstanceKey", 456),
-    ("zeebe.job.elementId", "service_task_123")
-])
+@validate_custom_parameters(
+    required_params=[
+        ("zeebe.job.key", 123),
+        ("zeebe.job.type", "testTask"),
+        ("zeebe.job.bpmnProcessId", "test_process"),
+        ("zeebe.job.processInstanceKey", 456),
+        ("zeebe.job.elementId", "service_task_123"),
+    ]
+)
 def test_execute_one_job():
     dummy_adapter = DummyZeebeAdapter()
 
     # Build a Job with fixed values
     job = Job(
         key=123,
-        type="testTask", # must match router.task(task_type="testTask")
+        type="testTask",  # must match router.task(task_type="testTask")
         bpmn_process_id="test_process",
         process_instance_key=456,
         process_definition_version=1,
@@ -52,7 +55,7 @@ def test_execute_one_job():
         retries=3,
         deadline=0,
         variables={"x": 33},
-        status=JobStatus.Running
+        status=JobStatus.Running,
     )
 
     # JobExecutor constructor params init
@@ -67,4 +70,4 @@ def test_execute_one_job():
     job_controller = JobController(job, dummy_adapter)
 
     asyncio.run(job_executor.execute_one_job(job, job_controller))
-    assert job.variables["x"] == 33
+    assert job.variables["x"] == 33