From dd4541ed9d7d557c2c395b1b716cea2cb51e1a44 Mon Sep 17 00:00:00 2001
From: Julien Barreau <julien.barreau@forestadmin.com>
Date: Fri, 20 Dec 2024 09:59:30 +0100
Subject: [PATCH] chore: fix and improve tests

---
 .../resources/test_capabilities_resource.py   | 51 +++++++++++++++++--
 src/agent_toolkit/tests/test_agent_toolkit.py |  6 ++-
 2 files changed, 53 insertions(+), 4 deletions(-)

diff --git a/src/agent_toolkit/tests/resources/test_capabilities_resource.py b/src/agent_toolkit/tests/resources/test_capabilities_resource.py
index 0ba6a52e..0e04db8d 100644
--- a/src/agent_toolkit/tests/resources/test_capabilities_resource.py
+++ b/src/agent_toolkit/tests/resources/test_capabilities_resource.py
@@ -17,6 +17,7 @@
 from forestadmin.agent_toolkit.services.permissions.permission_service import PermissionService
 from forestadmin.agent_toolkit.utils.context import Request, RequestMethod, Response, User
 from forestadmin.datasource_toolkit.collections import Collection
+from forestadmin.datasource_toolkit.datasource_customizer.datasource_customizer import DatasourceCustomizer
 from forestadmin.datasource_toolkit.datasources import Datasource
 from forestadmin.datasource_toolkit.interfaces.fields import FieldType, Operator, PrimitiveType
 
@@ -94,6 +95,12 @@ def setUpClass(cls) -> None:
                     "type": FieldType.COLUMN,
                     "filter_operators": set([]),
                 },
+                "created_at": {
+                    "column_type": PrimitiveType.DATE,
+                    "is_primary_key": False,
+                    "type": FieldType.COLUMN,
+                    "filter_operators": set([Operator.EQUAL, Operator.GREATER_THAN, Operator.LESS_THAN]),
+                },
             }  # type:ignore
         )
         cls.datasource.add_collection(cls.book_collection)
@@ -108,9 +115,17 @@ def setUpClass(cls) -> None:
             timezone=zoneinfo.ZoneInfo("Europe/Paris"),
             request={"ip": "127.0.0.1"},
         )
+        cls.datasource_customizer = DatasourceCustomizer()
+        cls.datasource_customizer.add_datasource(cls.datasource)
 
     def setUp(self) -> None:
-        self.capabilities_resource = CapabilitiesResource(self.datasource, self.ip_white_list_service, self.options)
+        self.capabilities_resource = CapabilitiesResource(
+            self.datasource_customizer.composite_datasource,
+            self.loop.run_until_complete(self.datasource_customizer.get_datasource()),
+            self.permission_service,
+            self.ip_white_list_service,
+            self.options,
+        )
 
     def test_dispatch_should_not_dispatch_to_capabilities_when_no_post_request(self):
         for method in [RequestMethod.DELETE, RequestMethod.OPTIONS, RequestMethod.PUT]:
@@ -127,7 +142,7 @@ def test_dispatch_should_not_dispatch_to_capabilities_when_no_post_request(self)
 
             self.assertEqual(response.status, 405)
 
-    def test_dispatch_should_return_correct_collection_and_fields_capabilities(self):
+    def test_dispatch_should_return_correct_collection_and_fields_capabilities_with_emulated_operators(self):
         request = Request(
             method=RequestMethod.POST,
             query={},
@@ -148,10 +163,40 @@ def test_dispatch_should_return_correct_collection_and_fields_capabilities(self)
                         {
                             "name": "id",
                             "type": "Number",
-                            "operators": ["equal", "greater_than", "in", "less_than"],
+                            "operators": ["blank", "equal", "greater_than", "in", "less_than", "missing"],
                         },
                         {"name": "name", "type": "String", "operators": []},
                         {"name": "cost", "type": "Number", "operators": []},
+                        {
+                            "name": "created_at",
+                            "type": "Date",
+                            "operators": [
+                                "after",
+                                "after_x_hours_ago",
+                                "before",
+                                "before_x_hours_ago",
+                                "blank",
+                                "equal",
+                                "future",
+                                "greater_than",
+                                "in",
+                                "less_than",
+                                "missing",
+                                "past",
+                                "previous_month",
+                                "previous_month_to_date",
+                                "previous_quarter",
+                                "previous_quarter_to_date",
+                                "previous_week",
+                                "previous_week_to_date",
+                                "previous_x_days",
+                                "previous_x_days_to_date",
+                                "previous_year",
+                                "previous_year_to_date",
+                                "today",
+                                "yesterday",
+                            ],
+                        },
                     ],
                 }
             ],
diff --git a/src/agent_toolkit/tests/test_agent_toolkit.py b/src/agent_toolkit/tests/test_agent_toolkit.py
index 2aceffeb..e3036b7b 100644
--- a/src/agent_toolkit/tests/test_agent_toolkit.py
+++ b/src/agent_toolkit/tests/test_agent_toolkit.py
@@ -100,7 +100,11 @@ def test_property_resources(
 
             mocked_authentication_resource.assert_called_once_with(agent._ip_white_list_service, agent.options)
             mocked_capabilities_resource.assert_called_once_with(
-                agent.customizer.composite_datasource, agent._ip_white_list_service, agent.options
+                agent.customizer.composite_datasource,
+                "fake_datasource",
+                agent._permission_service,
+                agent._ip_white_list_service,
+                agent.options,
             )
             mocked_crud_resource.assert_called_once_with(
                 agent.customizer.composite_datasource,