seitime-frappe/frappe/tests/classes/integration_test_case.py

import copy
import logging
from contextlib import AbstractContextManager, contextmanager

import frappe
from frappe.utils import cint

from ..utils.generators import make_test_records
from .unit_test_case import UnitTestCase

logger = logging.Logger(__file__)


class IntegrationTestCase(UnitTestCase):
	"""Integration test class for Frappe tests.

	Key features:
	- Automatic database setup and teardown
	- Utilities for managing database connections
	- Context managers for query counting and Redis call monitoring
	- Lazy loading of test record dependencies

	Note: If you override `setUpClass`, make sure to call `super().setUpClass()`
	to maintain the functionality of this base class.
	"""

	TEST_SITE = "test_site"

	SHOW_TRANSACTION_COMMIT_WARNINGS = False
	maxDiff = 10_000  # prints long diffs but useful in CI

	@classmethod
	def setUpClass(cls) -> None:
		super().setUpClass()

		# Site initialization
		cls.TEST_SITE = getattr(frappe.local, "site", None) or cls.TEST_SITE
		frappe.init(cls.TEST_SITE)
		cls.ADMIN_PASSWORD = frappe.get_conf(cls.TEST_SITE).admin_password

		cls._primary_connection = frappe.local.db
		cls._secondary_connection = None

		# Create test record dependencies
		cls._newly_created_test_records = []
		if cls.doctype and cls.doctype not in frappe.local.test_objects:
			cls._newly_created_test_records += make_test_records(cls.doctype)
		for doctype in getattr(cls.module, "test_dependencies", []):
			if doctype not in frappe.local.test_objects:
				cls._newly_created_test_records += make_test_records(doctype)

		# flush changes done so far to avoid flake
		frappe.db.commit()
		if cls.SHOW_TRANSACTION_COMMIT_WARNINGS:
			frappe.db.before_commit.add(_commit_watcher)

		# enqueue teardown actions (executed in LIFO order)
		cls.addClassCleanup(_restore_thread_locals, copy.deepcopy(frappe.local.flags))
		cls.addClassCleanup(_rollback_db)

	@classmethod
	def tearDownClass(cls) -> None:
		# Add any necessary teardown code here
		super().tearDownClass()

	def setUp(self) -> None:
		super().setUp()
		# Add any per-test setup code here

	def tearDown(self) -> None:
		# Add any per-test teardown code here
		super().tearDown()

	@contextmanager
	def primary_connection(self) -> AbstractContextManager[None]:
		"""Switch to primary DB connection

		This is used for simulating multiple users performing actions by simulating two DB connections"""
		try:
			current_conn = frappe.local.db
			frappe.local.db = self._primary_connection
			yield
		finally:
			frappe.local.db = current_conn

	@contextmanager
	def secondary_connection(self) -> AbstractContextManager[None]:
		"""Switch to secondary DB connection."""
		if self._secondary_connection is None:
			frappe.connect()  # get second connection
			self._secondary_connection = frappe.local.db

		try:
			current_conn = frappe.local.db
			frappe.local.db = self._secondary_connection
			yield
		finally:
			frappe.local.db = current_conn
			self.addCleanup(self._rollback_connections)

	def _rollback_connections(self) -> None:
		self._primary_connection.rollback()
		self._secondary_connection.rollback()

	@contextmanager
	def assertQueryCount(self, count: int) -> AbstractContextManager[None]:
		queries = []

		def _sql_with_count(*args, **kwargs):
			ret = orig_sql(*args, **kwargs)
			queries.append(args[0].last_query)
			return ret

		try:
			orig_sql = frappe.db.__class__.sql
			frappe.db.__class__.sql = _sql_with_count
			yield
			self.assertLessEqual(len(queries), count, msg="Queries executed: \n" + "\n\n".join(queries))
		finally:
			frappe.db.__class__.sql = orig_sql

	@contextmanager
	def assertRedisCallCounts(self, count: int) -> AbstractContextManager[None]:
		commands = []

		def execute_command_and_count(*args, **kwargs):
			ret = orig_execute(*args, **kwargs)
			key_len = 2
			if "H" in args[0]:
				key_len = 3
			commands.append((args)[:key_len])
			return ret

		try:
			orig_execute = frappe.cache.execute_command
			frappe.cache.execute_command = execute_command_and_count
			yield
			self.assertLessEqual(
				len(commands), count, msg="commands executed: \n" + "\n".join(str(c) for c in commands)
			)
		finally:
			frappe.cache.execute_command = orig_execute

	@contextmanager
	def assertRowsRead(self, count: int) -> AbstractContextManager[None]:
		rows_read = 0

		def _sql_with_count(*args, **kwargs):
			nonlocal rows_read

			ret = orig_sql(*args, **kwargs)
			# count of last touched rows as per DB-API 2.0 https://peps.python.org/pep-0249/#rowcount
			rows_read += cint(frappe.db._cursor.rowcount)
			return ret

		try:
			orig_sql = frappe.db.sql
			frappe.db.sql = _sql_with_count
			yield
			self.assertLessEqual(rows_read, count, msg="Queries read more rows than expected")
		finally:
			frappe.db.sql = orig_sql

	@contextmanager
	def switch_site(self, site: str) -> AbstractContextManager[None]:
		"""Switch connection to different site.
		Note: Drops current site connection completely."""

		try:
			old_site = frappe.local.site
			frappe.init(site, force=True)
			frappe.connect()
			yield
		finally:
			frappe.init(old_site, force=True)
			frappe.connect()

	@staticmethod
	@contextmanager
	def change_settings(doctype, settings_dict=None, /, commit=False, **settings):
		"""A context manager to ensure that settings are changed before running
		function and restored after running it regardless of exceptions occurred.
		This is useful in tests where you want to make changes in a function but
		don't retain those changes.
		import and use as decorator to cover full function or using `with` statement.

		example:
		@change_settings("Print Settings", {"send_print_as_pdf": 1})
		def test_case(self):
		        ...

		@change_settings("Print Settings", send_print_as_pdf=1)
		def test_case(self):
		        ...
		"""

		if settings_dict is None:
			settings_dict = settings

		try:
			settings = frappe.get_doc(doctype)
			# remember setting
			previous_settings = copy.deepcopy(settings_dict)
			for key in previous_settings:
				previous_settings[key] = getattr(settings, key)

			# change setting
			for key, value in settings_dict.items():
				setattr(settings, key, value)
			settings.save(ignore_permissions=True)
			# singles are cached by default, clear to avoid flake
			frappe.db.value_cache[settings] = {}
			if commit:
				frappe.db.commit()
			yield  # yield control to calling function

		finally:
			# restore settings
			settings = frappe.get_doc(doctype)
			for key, value in previous_settings.items():
				setattr(settings, key, value)
			settings.save(ignore_permissions=True)
			if commit:
				frappe.db.commit()


def _commit_watcher():
	import traceback

	logger.warning("Transaction committed during tests.")
	traceback.print_stack(limit=10)


def _rollback_db():
	frappe.db.value_cache = {}
	frappe.db.rollback()


def _restore_thread_locals(flags):
	frappe.local.flags = flags
	frappe.local.error_log = []
	frappe.local.message_log = []
	frappe.local.debug_log = []
	frappe.local.conf = frappe._dict(frappe.get_site_config())
	frappe.local.cache = {}
	frappe.local.lang = "en"
	frappe.local.preload_assets = {"style": [], "script": [], "icons": []}

	if hasattr(frappe.local, "request"):
		delattr(frappe.local, "request")