Refactor staking-related queries to separate file

test/functional/test_framework/mintlayer/staking.py

@@ -0,0 +1,42 @@
+import substrateinterface
+
+
+class Staking(object):
+    def __init__(self, substrate):
+        self.substrate = substrate
+
+    """ Query the node for the staking ledger """
+    def get_staking_ledger(self):
+        query = self.substrate.query_map(
+            module='Staking',
+            storage_function='Ledger'
+        )
+
+        return ((h, o.value) for (h, o) in query)
+
+    """ accesses current era """
+    def current_era(self):
+        query = self.substrate.query(
+            module='Staking',
+            storage_function='CurrentEra'
+        )
+        return query
+
+    """ gets the staking ledger of the given key """
+    def get_ledger(self, keypair):
+        query = self.substrate.query_map(
+            module='Staking',
+            storage_function='Ledger'
+        )
+
+        matching = lambda e: e[0].value == keypair.ss58_address
+
+        return filter(matching, ((h, o.value) for (h, o) in query))
+
+    """ returns what era for the user to be able to withdraw funds """
+    def withdrawal_era(self, keypair):
+        ledger = list(self.get_ledger(keypair))
+        if ledger:
+            return ledger[0][1]['unlocking'][0]['era']
+        else:
+            print("no funds to withdraw")

test/functional/test_framework/mintlayer/utxo.py

@@ -6,6 +6,7 @@
 import scalecodec
 import os
 import logging
+from staking import Staking
 
 """ Client. A thin wrapper over SubstrateInterface """
 class Client():
@@ -22,6 +23,8 @@ def __init__(self, url="ws://127.0.0.1", port=9944):
             type_registry=custom_type_registry
         )
 
+        self.staking = Staking(self.substrate)
+
     """ SCALE-encode given object in JSON format """
     def encode_obj(self, ty, obj):
         return self.substrate.encode_scale(ty, obj)
@@ -72,47 +75,21 @@ def get_staking_count(self, stash_keypair):
         return filter(matching , staking_count)
 
 
-    # TODO: move to a separate file
     """ accesses pallet-staking to retrieve the ledger """
     def get_staking_ledger(self):
-        query = self.substrate.query_map(
-            module='Staking',
-            storage_function='Ledger'
-        )
+        return self.staking.get_staking_ledger()
 
-        return ((h, o.value) for (h, o) in query)
-
-    # TODO: move to a separate file
     """ accesses current era """
     def current_era(self):
-        query = self.substrate.query(
-            module='Staking',
-            storage_function='CurrentEra'
-        )
-        return query
+        return self.staking.current_era()
 
-    # TODO: move to a separate file
     """ gets the staking ledger of the given key """
     def get_ledger(self, keypair):
-        query = self.substrate.query_map(
-            module='Staking',
-            storage_function='Ledger'
-        )
-
-        matching = lambda e: e[0].value == keypair.ss58_address
-
-        return filter(matching, ((h, o.value) for (h, o) in query))
+        return self.staking.get_ledger()
 
-    # TODO: move to a separate file
     """ returns what era for the user to be able to withdraw funds """
     def withdrawal_era(self, keypair):
-        ledger = list(self.get_ledger(keypair))
-        if ledger:
-            return ledger[0][1]['unlocking'][0]['era']
-        else:
-            print("no funds to withdraw")
-
-
+        return self.staking.withdrawal_era(keypair)
 
     """ Submit a transaction onto the blockchain """
     def submit(self, keypair, tx):