From eee7bf0ee2ca60d7815ed3b3827cdc0d754d5f6a Mon Sep 17 00:00:00 2001
From: nolash <dev@holbrook.no>
Date: Fri, 4 Jun 2021 19:37:12 +0200
Subject: [PATCH] Add tx count query

---
 chainqueue/sql/query.py | 15 +++++++++
 tests/base.py           |  1 +
 tests/test_query.py     | 70 ++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 85 insertions(+), 1 deletion(-)

diff --git a/chainqueue/sql/query.py b/chainqueue/sql/query.py
index fe2df21..478899a 100644
--- a/chainqueue/sql/query.py
+++ b/chainqueue/sql/query.py
@@ -339,3 +339,18 @@ def get_account_tx(chain_spec, address, as_sender=True, as_recipient=True, count
     SessionBase.release_session(session)
 
     return txs
+
+
+def count_tx(chain_spec, address=None, status=None, status_target=None, session=None):
+    session = SessionBase.bind_session(session)
+    q = session.query(Otx.id)
+    q = q.join(TxCache)
+    if status != None:
+        if status_target == None:
+            status_target = status
+        q = q.filter(Otx.status.op('&')(status)==status_target)
+    if address != None:
+        q = q.filter(TxCache.sender==address)
+    result = q.count()
+    SessionBase.release_session(session)
+    return result
diff --git a/tests/base.py b/tests/base.py
index 83b0de0..b1822c7 100644
--- a/tests/base.py
+++ b/tests/base.py
@@ -88,6 +88,7 @@ class TestTxBase(TestOtxBase):
     def setUp(self):
         super(TestTxBase, self).setUp()
         self.bob = add_0x(os.urandom(20).hex())
+        self.carol = add_0x(os.urandom(20).hex())
         self.foo_token = add_0x(os.urandom(20).hex())
         self.bar_token = add_0x(os.urandom(20).hex())
         self.from_value = 42
diff --git a/tests/test_query.py b/tests/test_query.py
index c4bb2bf..74d4c4d 100644
--- a/tests/test_query.py
+++ b/tests/test_query.py
@@ -12,7 +12,14 @@ from hexathon import (
 # local imports
 from chainqueue.sql.query import *
 from chainqueue.sql.tx import create
-from chainqueue.sql.state import set_waitforgas
+from chainqueue.sql.state import (
+        set_waitforgas,
+        set_ready,
+        set_reserved,
+        set_sent,
+        set_final,
+        )
+from chainqueue.enum import StatusBits
 
 # test imports
 from tests.base import TestTxBase
@@ -162,5 +169,66 @@ class TestTxQuery(TestTxBase):
         self.assertEqual(len(txs.keys()), 1)
 
 
+    def test_count(self):
+        for i in range(3):
+            tx_hash = add_0x(os.urandom(32).hex())
+            signed_tx = add_0x(os.urandom(128).hex())
+            create(
+                    self.chain_spec,
+                    i,
+                    self.alice,
+                    tx_hash,
+                    signed_tx,
+                    session=self.session,
+                    )
+            txc = TxCache(
+                tx_hash,
+                self.alice,
+                self.bob,
+                self.foo_token,
+                self.bar_token,
+                self.from_value,
+                self.to_value,
+                session=self.session,
+                )
+            self.session.add(txc)
+            set_ready(self.chain_spec, tx_hash, session=self.session)
+            set_reserved(self.chain_spec, tx_hash, session=self.session)
+            if i > 0:
+                set_sent(self.chain_spec, tx_hash, session=self.session)
+            if i == 2:
+                set_final(self.chain_spec, tx_hash, session=self.session)
+
+        tx_hash = add_0x(os.urandom(32).hex())
+        signed_tx = add_0x(os.urandom(128).hex())
+        create(
+                self.chain_spec,
+                i,
+                self.bob,
+                tx_hash,
+                signed_tx,
+                session=self.session,
+                )
+        txc = TxCache(
+            tx_hash,
+            self.bob,
+            self.carol,
+            self.foo_token,
+            self.bar_token,
+            self.from_value,
+            self.to_value,
+            session=self.session,
+            )
+
+        self.session.add(txc)
+        set_ready(self.chain_spec, tx_hash, session=self.session)
+        set_reserved(self.chain_spec, tx_hash, session=self.session)
+        set_sent(self.chain_spec, tx_hash, session=self.session)
+        self.session.commit()
+
+        self.assertEqual(count_tx(self.chain_spec, status=StatusBits.IN_NETWORK | StatusBits.FINAL, status_target=StatusBits.IN_NETWORK), 2)
+        self.assertEqual(count_tx(self.chain_spec, address=self.alice, status=StatusBits.IN_NETWORK | StatusBits.FINAL, status_target=StatusBits.IN_NETWORK), 1)
+
+
 if __name__ == '__main__':
     unittest.main()