From 76ab556f61c6e981cc702a6b246ddec149133dee Mon Sep 17 00:00:00 2001
From: Tom Owers <tom@posthog.com>
Date: Tue, 3 Oct 2023 15:57:23 +0100
Subject: [PATCH] feat: Added simple testing for trends formula ast (#17737)

Added simple testing for trends formula ast
---
 .../utils/test/test_formula_ast.py            | 53 +++++++++++++++++++
 1 file changed, 53 insertions(+)
 create mode 100644 posthog/hogql_queries/utils/test/test_formula_ast.py

diff --git a/posthog/hogql_queries/utils/test/test_formula_ast.py b/posthog/hogql_queries/utils/test/test_formula_ast.py
new file mode 100644
index 0000000000000..a10a54f71ca29
--- /dev/null
+++ b/posthog/hogql_queries/utils/test/test_formula_ast.py
@@ -0,0 +1,53 @@
+from posthog.hogql_queries.utils.formula_ast import FormulaAST
+from posthog.test.base import APIBaseTest
+
+
+class TestFormulaAST(APIBaseTest):
+    def _get_formula_ast(self) -> FormulaAST:
+        formula = FormulaAST(data=[[1, 2, 3, 4], [1, 2, 3, 4]])
+        return formula
+
+    def test_addition(self):
+        formula = self._get_formula_ast()
+        response = formula.call("A+1")
+        self.assertListEqual([2, 3, 4, 5], response)
+
+    def test_subtraction(self):
+        formula = self._get_formula_ast()
+        response = formula.call("A-1")
+        self.assertListEqual([0, 1, 2, 3], response)
+
+    def test_multiplication(self):
+        formula = self._get_formula_ast()
+        response = formula.call("A*2")
+        self.assertListEqual([2, 4, 6, 8], response)
+
+    def test_division(self):
+        formula = self._get_formula_ast()
+        response = formula.call("A/2")
+        self.assertListEqual([0.5, 1, 1.5, 2], response)
+
+    def test_modulo(self):
+        formula = self._get_formula_ast()
+        response = formula.call("A%2")
+        self.assertListEqual([1, 0, 1, 0], response)
+
+    def test_power(self):
+        formula = self._get_formula_ast()
+        response = formula.call("A**2")
+        self.assertListEqual([1, 4, 9, 16], response)
+
+    def test_constants(self):
+        formula = self._get_formula_ast()
+        response = formula.call("1")
+        self.assertListEqual([1, 1, 1, 1], response)
+
+    def test_named_values(self):
+        formula = self._get_formula_ast()
+        response = formula.call("A+B")
+        self.assertListEqual([2, 4, 6, 8], response)
+
+    def test_named_values_lower_case(self):
+        formula = self._get_formula_ast()
+        response = formula.call("a+b")
+        self.assertListEqual([2, 4, 6, 8], response)