Re-instate collection of WorkLock Refund events.

nucypher/utilities/prometheus/collector.py

@@ -33,10 +33,9 @@ from nucypher.blockchain.eth.registry import BaseContractRegistry
 from nucypher.datastore.datastore import RecordNotFound
 from nucypher.datastore.models import Workorder, PolicyArrangement
 
-from prometheus_client.metrics import MetricWrapperBase
 from prometheus_client.registry import CollectorRegistry
 
-from typing import Dict, List, Union
+from typing import Dict, Union
 
 ContractAgents = Union[StakingEscrowAgent, WorkLockAgent, PolicyManagerAgent]
 
@@ -432,3 +431,21 @@ class WorkerBondedEventMetricsCollector(EventMetricsCollector):
         super()._event_occurred(event)
         self.metrics["current_worker_is_me_gauge"].set(
             self.contract_agent.get_worker_from_staker(self.staker_address) == self.worker_address)
+
+
+class WorkLockRefundEventMetricsCollector(EventMetricsCollector):
+    """Collector for WorkLock Refund event."""
+
+    def __init__(self, staker_address: ChecksumAddress, event_name: str = 'Refund', *args, **kwargs):
+        super().__init__(event_name=event_name, argument_filters={'sender': staker_address}, *args, **kwargs)
+        self.staker_address = staker_address
+
+    def initialize(self, metrics_prefix: str, registry: CollectorRegistry) -> None:
+        super().initialize(metrics_prefix=metrics_prefix, registry=registry)
+        self.metrics["worklock_deposited_eth_gauge"] = Gauge(f'{metrics_prefix}_worklock_current_deposited_eth',
+                                                             'Worklock deposited ETH',
+                                                             registry=registry)
+
+    def _event_occurred(self, event) -> None:
+        super()._event_occurred(event)
+        self.metrics["worklock_deposited_eth_gauge"].set(self.contract_agent.get_deposited_eth(self.staker_address))

nucypher/utilities/prometheus/metrics.py

@@ -41,14 +41,15 @@ from nucypher.utilities.prometheus.collector import (
     ReStakeEventMetricsCollector,
     WindDownEventMetricsCollector,
     WorkerBondedEventMetricsCollector,
-    CommitmentMadeEventMetricsCollector)
+    CommitmentMadeEventMetricsCollector,
+    WorkLockRefundEventMetricsCollector)
 
 from typing import List
 
 from twisted.internet import reactor, task
 from twisted.web.resource import Resource
 
-from nucypher.blockchain.eth.agents import ContractAgency, StakingEscrowAgent, PolicyManagerAgent
+from nucypher.blockchain.eth.agents import ContractAgency, StakingEscrowAgent, PolicyManagerAgent, WorkLockAgent
 
 
 class PrometheusMetricsConfig:
@@ -205,6 +206,11 @@ def create_metrics_collectors(ursula: 'Ursula', metrics_prefix: str) -> List[Met
                                                                             metrics_prefix=metrics_prefix)
         collectors.extend(staking_events_collectors)
 
+        # WorkLock Events
+        worklock_events_collectors = create_worklock_events_metric_collectors(ursula=ursula,
+                                                                              metrics_prefix=metrics_prefix)
+        collectors.extend(worklock_events_collectors)
+
         # Policy Events
         policy_events_collectors = create_policy_events_metric_collectors(ursula=ursula,
                                                                           metrics_prefix=metrics_prefix)
@@ -283,6 +289,25 @@ def create_staking_events_metric_collectors(ursula: 'Ursula', metrics_prefix: st
     return collectors
 
 
+def create_worklock_events_metric_collectors(ursula: 'Ursula', metrics_prefix: str) -> List[MetricsCollector]:
+    """Create collectors for worklock-related events."""
+    collectors: List[MetricsCollector] = []
+    worklock_agent = ContractAgency.get_agent(WorkLockAgent, registry=ursula.registry)
+    staker_address = ursula.checksum_address
+
+    # Refund
+    collectors.append(WorkLockRefundEventMetricsCollector(
+        event_args_config={
+            "refundETH": (Gauge, f'{metrics_prefix}_worklock_refund_refundETH',
+                          'Refunded ETH'),
+        },
+        staker_address=staker_address,
+        contract_agent=worklock_agent,
+    ))
+
+    return collectors
+
+
 def create_policy_events_metric_collectors(ursula: 'Ursula', metrics_prefix: str) -> List[MetricsCollector]:
     """Create collectors for policy-related events."""
     collectors: List[MetricsCollector] = []