mercredi 4 mai 2016

Unable to raise OperationalError with Mock()

I'm writing a test for a bit of code that will execute a SQL query, and retry a specified amount of times if the query times out. If it doesn't succeed, I want it to set an alert in the object that will do something appropriate. My unittest is trying to capture this functionality, but I'm unable to raise the exception correctly. Here is my unit test:

def test_query_data_retries_the_correct_amount_of_times(self):
        shard = cu.Shard()
        fake_cursor = MagicMock()
        fake_cursor.execute = Mock(side_effect=ppg2.OperationalError('Operation timed out suckafoo'))
        def return_fake_cursor():
            return Mock(return_value=fake_cursor)
        shard._return_cursor = return_fake_cursor()
        self.assertRaises(ppg2.OperationalError, shard._query_data)
        self.assertEqual(fake_cursor.execute.call_count, 5)

Here is my function:

def _query_data(self, sql_query: str=None, retries: int=5):
        cursor = self._return_cursor(self.canvas_conn)
        success = False
        for attempt in range(retries):
            try:
                cursor.execute(sql_query)
            except (ppg2.ProgrammingError, ppg2.OperationalError) as e:
                if 'Operation timed out' in str(e):
                    print('Retrying, attempt {}'.format(attempt))
                    sleep(3)
            else:
                success = True
                # Only increment if successful.
                self.reporter.total_records += cursor.rowcount
            finally:
                if not success and self.historical_connect:
                    msg = 'Connected on last run, no longer available.'
                    self._set_for_alert(msg)
                    self.historical_connect = False
                return self._iterate_results(cursor)

Aucun commentaire:

Enregistrer un commentaire