From ba25420312b1102af34cd197d2c61e0ea3c0d93d Mon Sep 17 00:00:00 2001
From: Raphael Collet <rco@odoo.com>
Date: Tue, 12 Dec 2017 15:52:35 +0100
Subject: [PATCH] [REF] tests: decorators and method to count queries (#21499)

Add the context manager method `assertQueryCount` to count queries,
the decorator `@users(...)` to execute a test for several users, and
the decorator `@warmup` to execute a test twice (warm-up, then real).
---
 addons/test_mail/tests/test_performance.py    | 107 +++++++------
 .../tests/test_performance.py                 | 144 +++++-------------
 odoo/tests/common.py                          |  90 ++++++++++-
 3 files changed, 185 insertions(+), 156 deletions(-)

diff --git a/addons/test_mail/tests/test_performance.py b/addons/test_mail/tests/test_performance.py
index 06757e15dcf..7d178a5662e 100644
--- a/addons/test_mail/tests/test_performance.py
+++ b/addons/test_mail/tests/test_performance.py
@@ -1,10 +1,10 @@
 # -*- coding: utf-8 -*-
 # Part of Odoo. See LICENSE file for full copyright and licensing details.
 
-from odoo.addons.test_performance.tests.test_performance import TestPerformance, queryCount
+from odoo.tests.common import TransactionCase, users, warmup
 
 
-class TestMailPerformance(TestPerformance):
+class TestMailPerformance(TransactionCase):
 
     def setUp(self):
         super(TestMailPerformance, self).setUp()
@@ -21,75 +21,83 @@ class TestMailPerformance(TestPerformance):
             'groups_id': [(6, 0, [self.env.ref('base.group_user').id])],
         })
 
-        # for performances test
-        self.admin = self.env.user
-        self.admin.login = 'admin'
-
-    @queryCount(admin=3, demo=3)
+    @users('admin', 'demo')
+    @warmup
     def test_read_mail(self):
         """ Read records inheriting from 'mail.thread'. """
         records = self.env['test_performance.mail'].search([])
         self.assertEqual(len(records), 5)
-        self.resetQueryCount()
 
-        # without cache
-        for record in records:
-            record.partner_id.country_id.name
+        with self.assertQueryCount(admin=3, demo=3):
+            # without cache
+            for record in records:
+                record.partner_id.country_id.name
 
-        # with cache
-        for record in records:
-            record.partner_id.country_id.name
+        with self.assertQueryCount(0):
+            # with cache
+            for record in records:
+                record.partner_id.country_id.name
 
-        # value_pc must have been prefetched, too
-        for record in records:
-            record.value_pc
+        with self.assertQueryCount(0):
+            # value_pc must have been prefetched, too
+            for record in records:
+                record.value_pc
 
-    @queryCount(admin=4, demo=4)
+    @users('admin', 'demo')
+    @warmup
     def test_write_mail(self):
         """ Write records inheriting from 'mail.thread' (no recomputation). """
         records = self.env['test_performance.mail'].search([])
         self.assertEqual(len(records), 5)
-        self.resetQueryCount()
 
-        records.write({'name': self.str('X')})
+        with self.assertQueryCount(admin=4, demo=4):
+            records.write({'name': 'X'})
 
-    @queryCount(admin=6, demo=6)
+    @users('admin', 'demo')
+    @warmup
     def test_write_mail_with_recomputation(self):
         """ Write records inheriting from 'mail.thread' (with recomputation). """
         records = self.env['test_performance.mail'].search([])
         self.assertEqual(len(records), 5)
-        self.resetQueryCount()
 
-        records.write({'value': self.int(20)})
+        with self.assertQueryCount(admin=6, demo=6):
+            records.write({'value': 42})
 
-    @queryCount(admin=20, demo=31)
+    @users('admin', 'demo')
+    @warmup
     def test_write_mail_with_tracking(self):
         """ Write records inheriting from 'mail.thread' (with field tracking). """
         record = self.env['test_performance.mail'].search([], limit=1)
         self.assertEqual(len(record), 1)
-        self.resetQueryCount()
 
-        record.track = self.str('X')
+        with self.assertQueryCount(admin=20, demo=31):
+            record.track = 'X'
 
-    @queryCount(admin=3, demo=3)
+    @users('admin', 'demo')
+    @warmup
     def test_create_mail(self):
         """ Create records inheriting from 'mail.thread' (without field tracking). """
         model = self.env['test_performance.mail']
-        model.with_context(tracking_disable=True).create({'name': self.str('X')})
 
-    @queryCount(admin=38, demo=54)
+        with self.assertQueryCount(admin=3, demo=3):
+            model.with_context(tracking_disable=True).create({'name': 'X'})
+
+    @users('admin', 'demo')
+    @warmup
     def test_create_mail_with_tracking(self):
         """ Create records inheriting from 'mail.thread' (with field tracking). """
-        model = self.env['test_performance.mail']
-        model.create({'name': self.str('Y')})
+        with self.assertQueryCount(admin=38, demo=54):
+            self.env['test_performance.mail'].create({'name': 'X'})
 
-    @queryCount(admin=22, emp=29)
+    @users('admin', 'emp')
+    @warmup
     def test_simple(self):
         """ Create records inheriting from 'mail.thread' (simple models) """
-        self.env['mail.test.simple'].create({'name': self.str('Test')})
+        with self.assertQueryCount(admin=22, emp=29):
+            self.env['mail.test.simple'].create({'name': 'Test'})
 
 
-class TestAdvMailPerformance(TestPerformance):
+class TestAdvMailPerformance(TransactionCase):
 
     def setUp(self):
         super(TestAdvMailPerformance, self).setUp()
@@ -107,23 +115,24 @@ class TestAdvMailPerformance(TestPerformance):
             'groups_id': [(6, 0, [self.env.ref('base.group_user').id])],
         })
 
-        # for performances test
-        self.admin = self.env.user
-        self.admin.login = 'admin'
-
-    @queryCount(admin=25, emp=32)
+    @users('admin', 'emp')
+    @warmup
     def test_activity(self):
-        self.env['mail.test.activity'].create({'name': self.str('Test')})
+        model = self.env['mail.test.activity']
 
-    @queryCount(admin=46, emp=51)
+        with self.assertQueryCount(admin=25, emp=32):
+            model.create({'name': 'Test'})
+
+    @users('admin', 'emp')
+    @warmup
     def test_activity_full(self):
-        test_record_activity = self.env['mail.test.activity'].create({'name': self.str('Test')})
-        res_id = test_record_activity.id
-        self.resetQueryCount()
-
-        self.env['mail.activity'].with_context({
+        record = self.env['mail.test.activity'].create({'name': 'Test'})
+        model = self.env['mail.activity'].with_context({
             'default_res_model': 'mail.test.activity',
-        }).create({
-            'summary': self.str('Test Activity'),
-            'res_id': res_id,
         })
+
+        with self.assertQueryCount(admin=46, emp=51):
+            model.create({
+                'summary': 'Test Activity',
+                'res_id': record.id,
+            })
diff --git a/odoo/addons/test_performance/tests/test_performance.py b/odoo/addons/test_performance/tests/test_performance.py
index 005a68cc08e..3f9620159f6 100644
--- a/odoo/addons/test_performance/tests/test_performance.py
+++ b/odoo/addons/test_performance/tests/test_performance.py
@@ -1,138 +1,76 @@
 # -*- coding: utf-8 -*-
 # Part of Odoo. See LICENSE file for full copyright and licensing details.
 
-from contextlib import contextmanager
-import functools
-import logging
-
-from odoo.tests.common import TransactionCase
-
-_logger = logging.getLogger(__name__)
-sql_logger = logging.getLogger('odoo.sql_db')
-
-
-def queryCount(**counters):
-    """ Decorate a method to check the number of queries it makes. Counters
-    is a dict { 'user_login': expected_query_count } """
-    def decorate(func):
-        @functools.wraps(func)
-        def wrapper(self):
-            if hasattr(self, 'test_users'):
-                users = self.test_users.filtered(lambda user: user.login in counters.keys())
-            else:
-                users = self.env['res.users'].search([('login', 'in', list(counters.keys()))])
-            for user in users:
-                # switch user
-                self.uid = user.id
-                self.env = self.env(user=self.uid)
-                # warm up the caches
-                self._round = False
-                func(self)
-                self.env.cache.invalidate()
-                # test for real, and check query count
-                self._round = True
-                self.resetQueryCount()
-                func(self)
-                self.assertQueryCount(self.cr.sql_log_count - self._count,
-                                      counters[user.login], user.login)
-
-        return wrapper
-
-    return decorate
+from odoo.tests.common import TransactionCase, users, warmup
 
 
 class TestPerformance(TransactionCase):
 
-    def setUp(self):
-        super(TestPerformance, self).setUp()
-        self._round = False
-
-    def assertQueryCount(self, actual, expected, message):
-        self.assertLessEqual(actual, expected, message)
-        if actual < expected:
-            _logger.info("Warning: Got %d queries instead of %d: %s", actual, expected, message)
-
-    def resetQueryCount(self):
-        """ Reset the query counter. """
-        self._count = self.cr.sql_log_count
-
-    def str(self, value):
-        """ Return a value different from run to run. """
-        return value + 'z' if self._round else value
-
-    def int(self, value):
-        """ Return a value different from run to run. """
-        return value + 1 if self._round else value
-
-    @contextmanager
-    def logQueries(self):
-        """ Log the queries that are made in this scope. """
-        sql_log, level = self.cr.sql_log, sql_logger.getEffectiveLevel()
-        try:
-            sql_logger.setLevel(logging.DEBUG)
-            self.cr.sql_log = True
-            yield
-        finally:
-            self.cr.sql_log = sql_log
-            sql_logger.setLevel(level)
-
-    @queryCount(admin=3, demo=3)
+    @users('admin', 'demo')
+    @warmup
     def test_read_base(self):
         """ Read records. """
         records = self.env['test_performance.base'].search([])
         self.assertEqual(len(records), 5)
-        self.resetQueryCount()
 
-        # without cache
-        for record in records:
-            record.partner_id.country_id.name
+        with self.assertQueryCount(admin=3, demo=3):
+            # without cache
+            for record in records:
+                record.partner_id.country_id.name
 
-        # with cache
-        for record in records:
-            record.partner_id.country_id.name
+        with self.assertQueryCount(0):
+            # with cache
+            for record in records:
+                record.partner_id.country_id.name
 
-        # value_pc must have been prefetched, too
-        for record in records:
-            record.value_pc
+        with self.assertQueryCount(0):
+            # value_pc must have been prefetched, too
+            for record in records:
+                record.value_pc
 
-    @queryCount(admin=1, demo=1)
+    @users('admin', 'demo')
+    @warmup
     def test_write_base(self):
         """ Write records (no recomputation). """
         records = self.env['test_performance.base'].search([])
         self.assertEqual(len(records), 5)
-        self.resetQueryCount()
 
-        records.write({'name': self.str('X')})
+        with self.assertQueryCount(admin=1, demo=1):
+            records.write({'name': 'X'})
 
-    @queryCount(admin=3, demo=3)
+    @users('admin', 'demo')
+    @warmup
     def test_write_base_with_recomputation(self):
         """ Write records (with recomputation). """
         records = self.env['test_performance.base'].search([])
         self.assertEqual(len(records), 5)
-        self.resetQueryCount()
 
-        records.write({'value': self.int(20)})
+        with self.assertQueryCount(admin=3, demo=3):
+            records.write({'value': 42})
 
-    @queryCount(admin=6, demo=6)
+    @users('admin', 'demo')
+    @warmup
     def test_create_base(self):
         """ Create records. """
-        model = self.env['test_performance.base']
-        model.create({'name': self.str('X')})
+        with self.assertQueryCount(admin=6, demo=6):
+            self.env['test_performance.base'].create({'name': 'X'})
 
-    @queryCount(admin=38, demo=38)
+    @users('admin', 'demo')
+    @warmup
     def test_create_base_with_lines(self):
         """ Create records with one2many lines. """
-        model = self.env['test_performance.base']
-        model.create({
-            'name': self.str('Y'),
-            'line_ids': [(0, 0, {'value': val}) for val in range(10)],
-        })
+        with self.assertQueryCount(admin=38, demo=38):
+            self.env['test_performance.base'].create({
+                'name': 'X',
+                'line_ids': [(0, 0, {'value': val}) for val in range(10)],
+            })
 
-    @queryCount(admin=17, demo=17)
+    @users('admin', 'demo')
+    @warmup
     def test_create_base_with_tags(self):
         """ Create records with many2many tags. """
-        model = self.env['test_performance.base']
-        model.create({
-            'name': self.str('X'),
-            'tag_ids': [(0, 0, {'name': val}) for val in range(10)],
-        })
+        with self.assertQueryCount(admin=17, demo=17):
+            self.env['test_performance.base'].create({
+                'name': 'X',
+                'tag_ids': [(0, 0, {'name': val}) for val in range(10)],
+            })
diff --git a/odoo/tests/common.py b/odoo/tests/common.py
index 4e312da6d54..fcbd98b16ca 100644
--- a/odoo/tests/common.py
+++ b/odoo/tests/common.py
@@ -24,6 +24,7 @@ from datetime import datetime, timedelta, date
 from pprint import pformat
 
 import requests
+from decorator import decorator
 from lxml import etree, html
 
 from odoo.models import BaseModel
@@ -121,10 +122,21 @@ class BaseCase(TreeCase):
     """
 
     longMessage = True      # more verbose error message by default: https://www.odoo.com/r/Vmh
+    warm = True             # False during warm-up phase (see :func:`warmup`)
 
     def cursor(self):
         return self.registry.cursor()
 
+    @property
+    def uid(self):
+        """ Get the current uid. """
+        return self.env.uid
+
+    @uid.setter
+    def uid(self, user):
+        """ Set the uid by changing the test's environment. """
+        self.env = self.env(user=user)
+
     def ref(self, xid):
         """ Returns database ID for the provided :term:`external identifier`,
         shortcut for ``get_object_reference``
@@ -162,6 +174,35 @@ class BaseCase(TreeCase):
         else:
             return self._assertRaises(exception)
 
+    @contextmanager
+    def assertQueryCount(self, default=0, **counters):
+        """ Context manager that counts queries. It may be invoked either with
+            one value, or with a set of named arguments like ``login=value``::
+
+                with self.assertQueryCount(42):
+                    ...
+
+                with self.assertQueryCount(admin=3, demo=5):
+                    ...
+
+            The second form is convenient when used with :func:`users`.
+        """
+        if self.warm:
+            login = self.env.user.login
+            expected = counters.get(login, default)
+            count0 = self.cr.sql_log_count
+            yield
+            count = self.cr.sql_log_count - count0
+            if not count <= expected:
+                msg = "Query count for user %s: got %d instead of %d"
+                self.fail(msg % (login, count, expected))
+            elif count < expected:
+                logger = logging.getLogger(type(self).__module__)
+                msg = "Query count for user %s: got %d instead of %d"
+                logger.info(msg, login, count, expected)
+        else:
+            yield
+
     def shortDescription(self):
         doc = self._testMethodDoc
         return doc and ' '.join(l.strip() for l in doc.splitlines() if not l.isspace()) or None
@@ -182,9 +223,8 @@ class TransactionCase(BaseCase):
         self.registry = odoo.registry(get_db_name())
         #: current transaction's cursor
         self.cr = self.cursor()
-        self.uid = odoo.SUPERUSER_ID
         #: :class:`~odoo.api.Environment` for the current test case
-        self.env = api.Environment(self.cr, self.uid, {})
+        self.env = api.Environment(self.cr, odoo.SUPERUSER_ID, {})
 
         @self.addCleanup
         def reset():
@@ -219,8 +259,7 @@ class SingleTransactionCase(BaseCase):
         super(SingleTransactionCase, cls).setUpClass()
         cls.registry = odoo.registry(get_db_name())
         cls.cr = cls.registry.cursor()
-        cls.uid = odoo.SUPERUSER_ID
-        cls.env = api.Environment(cls.cr, cls.uid, {})
+        cls.env = api.Environment(cls.cr, odoo.SUPERUSER_ID, {})
 
     @classmethod
     def tearDownClass(cls):
@@ -468,6 +507,49 @@ class HttpCase(TransactionCase):
         cmd = ['phantomjs', phantomtest, json.dumps(options)]
         self.phantom_run(cmd, timeout)
 
+
+def users(*logins):
+    """ Decorate a method to execute it once for each given user. """
+    @decorator
+    def wrapper(func, *args, **kwargs):
+        self = args[0]
+        old_uid = self.uid
+        try:
+            # retrieve users
+            user_id = {
+                user.login: user.id
+                for user in self.env['res.users'].search([('login', 'in', list(logins))])
+            }
+            for login in logins:
+                # switch user
+                self.uid = user_id[login]
+                # execute func
+                func(*args, **kwargs)
+        finally:
+            self.uid = old_uid
+
+    return wrapper
+
+
+@decorator
+def warmup(func, *args, **kwargs):
+    """ Decorate a test method to run it twice: once for a warming up phase, and
+        a second time for real.  The test attribute ``warm`` is set to ``False``
+        during warm up, and ``True`` once the test is warmed up.  Note that the
+        effects of the warmup phase are rolled back thanks to a savepoint.
+    """
+    self = args[0]
+    # run once to warm up the caches
+    self.warm = False
+    self.cr.execute('SAVEPOINT test_warmup')
+    func(*args, **kwargs)
+    self.cr.execute('ROLLBACK TO SAVEPOINT test_warmup')
+    self.env.cache.invalidate()
+    # run once for real
+    self.warm = True
+    func(*args, **kwargs)
+
+
 def can_import(module):
     """ Checks if <module> can be imported, returns ``True`` if it can be,
     ``False`` otherwise.