1  """Test cases for helper_sql"""
  2  
  3  import unittest
  4  
  5  try:
  6      # Python 3
  7      from unittest.mock import patch
  8  except ImportError:
  9      # Python 2
 10      from mock import patch
 11  
 12  import pybitmessage.helper_sql as helper_sql
 13  
 14  
 15  class TestHelperSql(unittest.TestCase):
 16      """Test class for helper_sql"""
 17  
 18      @classmethod
 19      def setUpClass(cls):
 20          helper_sql.sql_available = True
 21  
 22      @patch("pybitmessage.helper_sql.sqlSubmitQueue.put")
 23      @patch("pybitmessage.helper_sql.sqlReturnQueue.get")
 24      def test_sqlquery_no_args(self, mock_sqlreturnqueue_get, mock_sqlsubmitqueue_put):
 25          """Test sqlQuery with no additional arguments"""
 26          mock_sqlreturnqueue_get.return_value = ("dummy_result", None)
 27          result = helper_sql.sqlQuery(
 28              "SELECT msgid FROM inbox where folder='inbox' ORDER BY received"
 29          )
 30          self.assertEqual(mock_sqlsubmitqueue_put.call_count, 2)
 31          self.assertEqual(result, "dummy_result")
 32  
 33      @patch("pybitmessage.helper_sql.sqlSubmitQueue.put")
 34      @patch("pybitmessage.helper_sql.sqlReturnQueue.get")
 35      def test_sqlquery_with_args(self, mock_sqlreturnqueue_get, mock_sqlsubmitqueue_put):
 36          """Test sqlQuery with additional arguments"""
 37          mock_sqlreturnqueue_get.return_value = ("dummy_result", None)
 38          result = helper_sql.sqlQuery(
 39              "SELECT address FROM addressbook WHERE address=?", "PB-5yfds868gbkj"
 40          )
 41          self.assertEqual(mock_sqlsubmitqueue_put.call_count, 2)
 42          self.assertEqual(result, "dummy_result")
 43  
 44      @patch("pybitmessage.helper_sql.sqlSubmitQueue.put")
 45      @patch("pybitmessage.helper_sql.sqlReturnQueue.get")
 46      def test_sqlexecute(self, mock_sqlreturnqueue_get, mock_sqlsubmitqueue_put):
 47          """Test sqlExecute with valid arguments"""
 48          mock_sqlreturnqueue_get.return_value = (None, 1)
 49          rowcount = helper_sql.sqlExecute(
 50              "UPDATE sent SET status = 'msgqueued'"
 51              "WHERE ackdata = ? AND folder = 'sent'",
 52              "1710652313",
 53          )
 54          self.assertEqual(mock_sqlsubmitqueue_put.call_count, 3)
 55          self.assertEqual(rowcount, 1)
 56  
 57      @patch("pybitmessage.helper_sql.SqlBulkExecute.execute")
 58      def test_sqlexecute_script(self, mock_execute):
 59          """Test sqlExecuteScript with a SQL script"""
 60          helper_sql.sqlExecuteScript(
 61              "CREATE TABLE test (id INTEGER); INSERT INTO test VALUES (1);"
 62          )
 63          self.assertTrue(mock_execute.assert_called)
 64  
 65      @patch("pybitmessage.helper_sql.sqlSubmitQueue.put")
 66      @patch(
 67          "pybitmessage.helper_sql.sqlReturnQueue.get",
 68      )
 69      def test_sqlexecute_chunked(self, mock_sqlreturnqueue_get, mock_sqlsubmitqueue_put):
 70          """Test sqlExecuteChunked with valid arguments"""
 71          # side_effect is list of return value (_, rowcount)
 72          # of sqlReturnQueue.get for each chunk
 73          CHUNK_COUNT = 6
 74          CHUNK_SIZE = 999
 75          ID_COUNT = CHUNK_COUNT * CHUNK_SIZE
 76          CHUNKS_ROWCOUNT_LIST = [50, 29, 28, 18, 678, 900]
 77          TOTAL_ROW_COUNT = sum(CHUNKS_ROWCOUNT_LIST)
 78          mock_sqlreturnqueue_get.side_effect = [(None, rowcount) for rowcount in CHUNKS_ROWCOUNT_LIST]
 79          args = []
 80          for i in range(0, ID_COUNT):
 81              args.append("arg{}".format(i))
 82          total_row_count_return = helper_sql.sqlExecuteChunked(
 83              "INSERT INTO table VALUES {}", ID_COUNT, *args
 84          )
 85          self.assertEqual(TOTAL_ROW_COUNT, total_row_count_return)
 86          self.assertTrue(mock_sqlsubmitqueue_put.called)
 87          self.assertTrue(mock_sqlreturnqueue_get.called)
 88  
 89      @patch("pybitmessage.helper_sql.sqlSubmitQueue.put")
 90      @patch("pybitmessage.helper_sql.sqlReturnQueue.get")
 91      def test_sqlexecute_chunked_with_idcount_zero(
 92          self, mock_sqlreturnqueue_get, mock_sqlsubmitqueue_put
 93      ):
 94          """Test sqlExecuteChunked with id count 0"""
 95          ID_COUNT = 0
 96          args = list()
 97          for i in range(0, ID_COUNT):
 98              args.append("arg{}".format(i))
 99          total_row_count = helper_sql.sqlExecuteChunked(
100              "INSERT INTO table VALUES {}", ID_COUNT, *args
101          )
102          self.assertEqual(total_row_count, 0)
103          self.assertFalse(mock_sqlsubmitqueue_put.called)
104          self.assertFalse(mock_sqlreturnqueue_get.called)
105  
106      @patch("pybitmessage.helper_sql.sqlSubmitQueue.put")
107      @patch("pybitmessage.helper_sql.sqlReturnQueue.get")
108      def test_sqlexecute_chunked_with_args_less(
109          self, mock_sqlreturnqueue_get, mock_sqlsubmitqueue_put
110      ):
111          """Test sqlExecuteChunked with length of args less than idcount"""
112          ID_COUNT = 12
113          args = ["args0", "arg1"]
114          total_row_count = helper_sql.sqlExecuteChunked(
115              "INSERT INTO table VALUES {}", ID_COUNT, *args
116          )
117          self.assertEqual(total_row_count, 0)
118          self.assertFalse(mock_sqlsubmitqueue_put.called)
119          self.assertFalse(mock_sqlreturnqueue_get.called)
120  
121      @patch("pybitmessage.helper_sql.sqlSubmitQueue.put")
122      @patch("pybitmessage.helper_sql.sqlSubmitQueue.task_done")
123      def test_sqlstored_procedure(self, mock_task_done, mock_sqlsubmitqueue_put):
124          """Test sqlStoredProcedure with a stored procedure name"""
125          helper_sql.sqlStoredProcedure("exit")
126          self.assertTrue(mock_task_done.called_once)
127          mock_sqlsubmitqueue_put.assert_called_with("terminate")
128  
129      @classmethod
130      def tearDownClass(cls):
131          helper_sql.sql_available = False