diff --git a/posthog/hogql/test/_test_parser.py b/posthog/hogql/test/_test_parser.py index fdd5dd946a2d12..765d4fbaab4de5 100644 --- a/posthog/hogql/test/_test_parser.py +++ b/posthog/hogql/test/_test_parser.py @@ -6,11 +6,18 @@ from posthog.hogql.errors import HogQLException, SyntaxException from posthog.hogql.parser import parse_expr, parse_order_expr, parse_select from posthog.hogql.visitor import clear_locations -from posthog.test.base import BaseTest +from posthog.test.base import BaseTest, MemoryLeakTestMixin def parser_test_factory(backend: Literal["python", "cpp"]): - class TestParser(BaseTest): + base_classes = (MemoryLeakTestMixin, BaseTest) if backend == "cpp" else (BaseTest,) + + class TestParser(*base_classes): + MEMORY_INCREASE_PER_PARSE_LIMIT_B = 10_000 + MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT = 0.1 + MEMORY_PRIMING_RUNS_N = 2 + MEMORY_LEAK_CHECK_RUNS_N = 100 + maxDiff = None def _expr(self, expr: str, placeholders: Optional[Dict[str, ast.Expr]] = None) -> ast.Expr: diff --git a/posthog/test/base.py b/posthog/test/base.py index 5457bbe4056bc8..17b52736e52a39 100644 --- a/posthog/test/base.py +++ b/posthog/test/base.py @@ -1,11 +1,12 @@ import datetime as dt import inspect import re +import resource import threading import uuid from contextlib import contextmanager from functools import wraps -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union from unittest.mock import patch import freezegun @@ -133,7 +134,7 @@ def validation_error_response( return {"type": "validation_error", "code": code, "detail": message, "attr": attr} -class TestMixin: +class TestMixin(TestCase if TYPE_CHECKING else object): CONFIG_ORGANIZATION_NAME: str = "Test" CONFIG_EMAIL: Optional[str] = "user1@posthog.com" CONFIG_PASSWORD: Optional[str] = "testpassword12345" @@ -194,6 +195,41 @@ def validate_basic_html(self, html_message, site_url, preheader=None): self.assertIn(preheader, html_message) # type: ignore +class MemoryLeakTestMixin(TestCase if TYPE_CHECKING else object): + MEMORY_INCREASE_PER_PARSE_LIMIT_B: int + """Parsing more than once can never increase memory by this much (on average)""" + MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT: float + """Parsing more than once can never increase memory by more than this factor * first run's increase (on average)""" + MEMORY_PRIMING_RUNS_N: int + """How many times to run every test method to prime the heap""" + MEMORY_LEAK_CHECK_RUNS_N: int + """How many times to run every test method to check for memory leaks""" + + def _callTestMethod(self, method): + mem_original_b = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + for _ in range(self.MEMORY_PRIMING_RUNS_N): # Priming runs + method() + mem_primed_b = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + for _ in range(self.MEMORY_LEAK_CHECK_RUNS_N): # Memory leak check runs + method() + mem_tested_b = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss + avg_memory_priming_increase_b = (mem_primed_b - mem_original_b) / self.MEMORY_PRIMING_RUNS_N + avg_memory_test_increase_b = (mem_tested_b - mem_primed_b) / self.MEMORY_LEAK_CHECK_RUNS_N + avg_memory_increase_factor = ( + avg_memory_test_increase_b / avg_memory_priming_increase_b if avg_memory_priming_increase_b else 0 + ) + self.assertLessEqual( + avg_memory_test_increase_b, + self.MEMORY_INCREASE_PER_PARSE_LIMIT_B, + f"Possible memory leak - exceeded {self.MEMORY_INCREASE_PER_PARSE_LIMIT_B}-byte limit of incremental memory per parse", + ) + self.assertLessEqual( + avg_memory_increase_factor, + self.MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT, + f"Possible memory leak - exceeded {self.MEMORY_INCREASE_INCREMENTAL_FACTOR_LIMIT*100:.2f}% limit of incremental memory per parse", + ) + + class BaseTest(TestMixin, ErrorResponsesMixin, TestCase): """ Base class for performing Postgres-based backend unit tests on.