diff --git a/examples/rapidpro/contacts.py b/examples/rapidpro/contacts.py index dd060c3..18414e9 100644 --- a/examples/rapidpro/contacts.py +++ b/examples/rapidpro/contacts.py @@ -1,7 +1,7 @@ from api.rapidpro import pyRapid contacts = pyRapid().contacts.get_contacts( - before="2023-01-02 00:00:00", after="2023-01-01 00:00:00" + end_datetime="2023-01-01 01:00:50", start_datetime="2023-01-01 00:00:00" ) -contacts.head(5) +print(contacts.collect()) diff --git a/examples/rapidpro/fields.py b/examples/rapidpro/fields.py index 6d68d51..93377d1 100644 --- a/examples/rapidpro/fields.py +++ b/examples/rapidpro/fields.py @@ -2,4 +2,4 @@ fields = pyRapid().fields.get_fields() -fields.head(5) +print(fields.collect()) diff --git a/examples/rapidpro/flows.py b/examples/rapidpro/flows.py index a11f221..d4059db 100644 --- a/examples/rapidpro/flows.py +++ b/examples/rapidpro/flows.py @@ -2,4 +2,4 @@ flows = pyRapid().flows.get_flows() -flows.head(5) +print(flows.collect()) diff --git a/examples/rapidpro/flowstarts.py b/examples/rapidpro/flowstarts.py index 847a1db..4321723 100644 --- a/examples/rapidpro/flowstarts.py +++ b/examples/rapidpro/flowstarts.py @@ -1,7 +1,7 @@ from api.rapidpro import pyRapid flowstarts = pyRapid().flow_starts.get_flowstarts( - before="2023-01-02 00:00:00", after="2023-01-01 00:00:00" + end_datetime="2023-01-02 00:00:00", start_datetime="2023-01-01 00:00:00" ) -flowstarts.head(5) +print(flowstarts.collect()) diff --git a/examples/rapidpro/groups.py b/examples/rapidpro/groups.py index 9a3dce9..99fbdec 100644 --- a/examples/rapidpro/groups.py +++ b/examples/rapidpro/groups.py @@ -2,4 +2,4 @@ groups = pyRapid().groups.get_groups() -groups.head(5) +print(groups.collect()) diff --git a/examples/rapidpro/runs.py b/examples/rapidpro/runs.py index f17b7fa..dcc7bf4 100644 --- a/examples/rapidpro/runs.py +++ b/examples/rapidpro/runs.py @@ -1,7 +1,7 @@ from api.rapidpro import pyRapid runs = pyRapid().runs.get_runs( - before="2024-10-01 01:30:00", after="2024-10-01 01:00:00" + end_datetime="2024-06-22 00:00:10", start_datetime="2024-06-22 00:00:00" ) -runs.head(5) +print(runs.collect()) diff --git a/pyproject.toml b/pyproject.toml index 3a391f9..663ed98 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,9 +33,15 @@ dev = [ "boto3-stubs>=1.37.19", "starlette>=0.46.1", "pytest-coverage>=0.0", + "pytest-mock>=3.14.1", "werkzeug>=3.1.3", ] +[tool.coverage.run] +omit = [ + "tests/*" +] + [build-system] requires = ["hatchling"] build-backend = "hatchling.build" diff --git a/rdw_ingestion_tools/api/__init__.py b/rdw_ingestion_tools/api/__init__.py index 8b9240c..3eb849c 100644 --- a/rdw_ingestion_tools/api/__init__.py +++ b/rdw_ingestion_tools/api/__init__.py @@ -2,7 +2,15 @@ from collections.abc import Iterator from pandas import DataFrame -from polars import LazyFrame, concat, json_normalize +from pandas import json_normalize as pd_json_normalize +from polars import ( + LazyFrame, + Object, + String, + col, + concat, + json_normalize, +) # https://github.com/astral-sh/ruff/issues/3388 from typing_extensions import Never # noqa: UP035 @@ -31,14 +39,14 @@ def concatenate( """ try: - df = concat([json_normalize(obj, sep="_") for obj in objs]) + df = concat([pd_json_normalize(obj, sep="_") for obj in objs]) except ValueError: df = DataFrame() return df -def concatenate_to_lf( +def concatenate_to_lazyframe( objs: list[dict] | dict[Never, Never] | list[Never] | Iterator, schema: dict ) -> LazyFrame: """ @@ -57,3 +65,46 @@ def concatenate_to_lf( lf = LazyFrame(schema=schema) return lf + + +def get_polars_schema( + object_columns: list[str], data: list[dict[str, Object]] +) -> dict[str, Object]: + """ + Creates a normalised LazyFrame and uses the schema to generate a schema + dictionary using the column names. + Columns that are `list` types need to be type `Object` before they can be cast + to string. + All other column types can be cast directly to string using the schema generated. + """ + # Create a dataframe to use to build the schema + columns = ( + json_normalize(data, separator="_", infer_schema_length=None) + .lazy() + .collect_schema() + .names() + ) + schema = { + column: (Object if column in object_columns else String) for column in columns + } + + return schema + + +def concatenate_to_string_lazyframe( + objs: list[dict] | dict[Never, Never] | list[Never] | Iterator, + object_columns: list[str], +) -> LazyFrame: + """ + Flattens JSON data. Returns a LazyFrame with String columns. + """ + data = list(objs) + + schema = get_polars_schema(data=data, object_columns=object_columns) + lf = ( + json_normalize(data, separator="_", schema=schema) + .lazy() + .with_columns(col(Object).map_elements(lambda x: str(x), return_dtype=String)) + ) + + return lf diff --git a/rdw_ingestion_tools/api/rapidpro/extensions/httpx.py b/rdw_ingestion_tools/api/rapidpro/extensions/httpx.py index 6217147..43bc5a3 100644 --- a/rdw_ingestion_tools/api/rapidpro/extensions/httpx.py +++ b/rdw_ingestion_tools/api/rapidpro/extensions/httpx.py @@ -1,4 +1,5 @@ from collections.abc import Iterator +from urllib.parse import unquote from httpx import Client @@ -14,9 +15,10 @@ def get_paginated( pages are left to return. """ + params = {**kwargs} while True: - response = client.get(url, params={**kwargs}) + response = client.get(url, params=params) response.raise_for_status() data: dict = response.json() @@ -25,6 +27,8 @@ def get_paginated( yield from results try: - url = data["next"].split("/v2/")[1] + cursor = data["next"].split("cursor=")[1].split("&")[0] + decoded_cursor = unquote(cursor) + params["cursor"] = decoded_cursor except AttributeError: break diff --git a/rdw_ingestion_tools/api/rapidpro/requests/contacts.py b/rdw_ingestion_tools/api/rapidpro/requests/contacts.py index 4acb16b..729a06e 100644 --- a/rdw_ingestion_tools/api/rapidpro/requests/contacts.py +++ b/rdw_ingestion_tools/api/rapidpro/requests/contacts.py @@ -1,8 +1,8 @@ from attrs import define from httpx import Client -from pandas import DataFrame +from polars import LazyFrame -from api import concatenate +from api import concatenate_to_string_lazyframe from ..extensions.httpx import get_paginated @@ -13,22 +13,28 @@ class Contacts: client: Client - def get_contacts(self, **kwargs: str | int) -> DataFrame: - """Get a pandas DataFrame of Rapidpro contacts. + def get_contacts( + self, start_datetime: str, end_datetime: str, **kwargs: str | int + ) -> LazyFrame: + """Get a Polars LazyFrame of Rapidpro contacts. This endpoint supports time-based filtering that allows to fetch results between two date parameters. Example: pyRapid.contacts.get_contacts( - before="2023-01-02T00:00:00", - after="2023-01-01T00:00:00" + end_datetime="2023-01-02T00:00:00", + start_datetime="2023-01-01T00:00:00" ) """ url = "contacts.json" - contacts_generator = get_paginated(self.client, url, **kwargs) + contacts_generator = get_paginated( + self.client, url, after=start_datetime, before=end_datetime, **kwargs + ) - contacts = concatenate(contacts_generator) + contacts = concatenate_to_string_lazyframe( + contacts_generator, object_columns=["urns", "groups"] + ) return contacts diff --git a/rdw_ingestion_tools/api/rapidpro/requests/fields.py b/rdw_ingestion_tools/api/rapidpro/requests/fields.py index 2966e43..3991a4f 100644 --- a/rdw_ingestion_tools/api/rapidpro/requests/fields.py +++ b/rdw_ingestion_tools/api/rapidpro/requests/fields.py @@ -1,8 +1,8 @@ from attrs import define from httpx import Client -from pandas import DataFrame +from polars import LazyFrame -from api import concatenate +from api import concatenate_to_string_lazyframe from ..extensions.httpx import get_paginated @@ -13,8 +13,8 @@ class Fields: client: Client - def get_fields(self, **kwargs: str | int) -> DataFrame: - """Get a pandas DataFrame of Rapidpro fields. + def get_fields(self, **kwargs: str | int) -> LazyFrame: + """Get a polars LazyFrame of Rapidpro fields. This endpoint does not support time-based filtering and can be called as: @@ -26,6 +26,6 @@ def get_fields(self, **kwargs: str | int) -> DataFrame: fields_generator = get_paginated(self.client, url, **kwargs) - fields = concatenate(fields_generator) + fields = concatenate_to_string_lazyframe(fields_generator, object_columns=[]) return fields diff --git a/rdw_ingestion_tools/api/rapidpro/requests/flow_starts.py b/rdw_ingestion_tools/api/rapidpro/requests/flow_starts.py index a3062f4..9bd0902 100644 --- a/rdw_ingestion_tools/api/rapidpro/requests/flow_starts.py +++ b/rdw_ingestion_tools/api/rapidpro/requests/flow_starts.py @@ -1,8 +1,8 @@ from attrs import define from httpx import Client -from pandas import DataFrame +from polars import LazyFrame -from api import concatenate +from api import concatenate_to_string_lazyframe from ..extensions.httpx import get_paginated @@ -13,22 +13,28 @@ class FlowStarts: client: Client - def get_flowstarts(self, **kwargs: str | int) -> DataFrame: - """Get a pandas DataFrame of Rapidpro flowstarts. + def get_flowstarts( + self, start_datetime: str, end_datetime: str, **kwargs: str | int + ) -> LazyFrame: + """Get a Polars LazyFrame of Rapidpro flowstarts. This endpoint supports time-based filtering that allows to fetch results between two date parameters. Example: pyRapid.flowstarts.get_flowstarts( - before="2023-01-02T00:00:00", - after="2023-01-01T00:00:00" + end_datetime="2023-01-02T00:00:00", + start_datetime="2023-01-01T00:00:00" ) """ url = "flow_starts.json" - flowstarts_generator = get_paginated(self.client, url, **kwargs) + flowstarts_generator = get_paginated( + self.client, url, after=start_datetime, before=end_datetime, **kwargs + ) - flowstarts = concatenate(flowstarts_generator) + flowstarts = concatenate_to_string_lazyframe( + flowstarts_generator, object_columns=["groups", "contacts"] + ) return flowstarts diff --git a/rdw_ingestion_tools/api/rapidpro/requests/flows.py b/rdw_ingestion_tools/api/rapidpro/requests/flows.py index f702fd4..af35f6f 100644 --- a/rdw_ingestion_tools/api/rapidpro/requests/flows.py +++ b/rdw_ingestion_tools/api/rapidpro/requests/flows.py @@ -1,8 +1,8 @@ from attrs import define from httpx import Client -from pandas import DataFrame +from polars import LazyFrame -from api import concatenate +from api import concatenate_to_string_lazyframe from ..extensions.httpx import get_paginated @@ -13,8 +13,8 @@ class Flows: client: Client - def get_flows(self, **kwargs: str | int) -> DataFrame: - """Get a pandas DataFrame of Rapidpro flows. + def get_flows(self, **kwargs: str | int) -> LazyFrame: + """Get a Polars LazyFrame of Rapidpro flows. This endpoint does not support time-based filtering and can be called as: @@ -26,6 +26,8 @@ def get_flows(self, **kwargs: str | int) -> DataFrame: flows_generator = get_paginated(self.client, url, **kwargs) - flows = concatenate(flows_generator) + flows = concatenate_to_string_lazyframe( + flows_generator, object_columns=["labels", "results", "parent_refs"] + ) return flows diff --git a/rdw_ingestion_tools/api/rapidpro/requests/groups.py b/rdw_ingestion_tools/api/rapidpro/requests/groups.py index 9a85ea9..6f0f732 100644 --- a/rdw_ingestion_tools/api/rapidpro/requests/groups.py +++ b/rdw_ingestion_tools/api/rapidpro/requests/groups.py @@ -1,8 +1,8 @@ from attrs import define from httpx import Client -from pandas import DataFrame +from polars import LazyFrame -from api import concatenate +from api import concatenate_to_string_lazyframe from ..extensions.httpx import get_paginated @@ -13,8 +13,8 @@ class Groups: client: Client - def get_groups(self, **kwargs: str | int) -> DataFrame: - """Get a pandas DataFrame of Rapidpro groups. + def get_groups(self, **kwargs: str | int) -> LazyFrame: + """Get a Polars LazyFrame of Rapidpro groups. This endpoint does not support time-based filtering and can be called as: @@ -26,6 +26,6 @@ def get_groups(self, **kwargs: str | int) -> DataFrame: groups_generator = get_paginated(self.client, url, **kwargs) - groups = concatenate(groups_generator) + groups = concatenate_to_string_lazyframe(groups_generator, object_columns=[]) return groups diff --git a/rdw_ingestion_tools/api/rapidpro/requests/runs.py b/rdw_ingestion_tools/api/rapidpro/requests/runs.py index 251e1d4..249336b 100644 --- a/rdw_ingestion_tools/api/rapidpro/requests/runs.py +++ b/rdw_ingestion_tools/api/rapidpro/requests/runs.py @@ -1,8 +1,8 @@ from attrs import define from httpx import Client -from pandas import DataFrame +from polars import LazyFrame -from api import concatenate +from api import concatenate_to_string_lazyframe from ..extensions.httpx import get_paginated @@ -13,22 +13,28 @@ class Runs: client: Client - def get_runs(self, **kwargs: str | int) -> DataFrame: - """Get a pandas DataFrame of Rapidpro runs. + def get_runs( + self, start_datetime: str, end_datetime: str, **kwargs: str | int + ) -> LazyFrame: + """Get a Polars LazyFrame of Rapidpro runs. This endpoint supports time-based filtering that allows to fetch results between two date parameters. Example: pyRapid.runs.get_runs( - before="2023-01-02T00:00:00", - after="2023-01-01T00:00:00" + end_datetime="2023-01-02T00:00:00", + start_datetime="2023-01-01T00:00:00" ) """ url = "runs.json" - runs_generator = get_paginated(self.client, url, **kwargs) + runs_generator = get_paginated( + self.client, url, after=start_datetime, before=end_datetime, **kwargs + ) - runs = concatenate(runs_generator) + runs = concatenate_to_string_lazyframe( + objs=runs_generator, object_columns=["path", "flow", "contact", "events"] + ) return runs diff --git a/rdw_ingestion_tools/api/turn_bq/requests/cards.py b/rdw_ingestion_tools/api/turn_bq/requests/cards.py index 1047865..1bafb32 100644 --- a/rdw_ingestion_tools/api/turn_bq/requests/cards.py +++ b/rdw_ingestion_tools/api/turn_bq/requests/cards.py @@ -1,4 +1,4 @@ -from api import concatenate_to_lf +from api import concatenate_to_lazyframe from attrs import define from httpx import Client from polars import LazyFrame @@ -14,23 +14,23 @@ class Cards: client: Client def get_cards(self) -> LazyFrame: - """Returns a pandas DataFrame of Turn Cards""" + """Returns a Polars LazyFrame of Turn Cards""" url = "cards/" cards_generator = get_paginated(self.client, url) - cards = concatenate_to_lf(cards_generator, cards_schema) + cards = concatenate_to_lazyframe(cards_generator, cards_schema) return cards def get_cards_by_id(self, card_id: int) -> LazyFrame: - """Returns a pandas DataFrame of Turn Cards by card_id""" + """Returns a Polars LazyFrame of Turn Cards by card_id""" url = f"cards/{card_id}" cards_generator = get_paginated(self.client, url) - cards = concatenate_to_lf(cards_generator, cards_schema) + cards = concatenate_to_lazyframe(cards_generator, cards_schema) return cards diff --git a/rdw_ingestion_tools/api/turn_bq/requests/chats.py b/rdw_ingestion_tools/api/turn_bq/requests/chats.py index b942782..7b4efef 100644 --- a/rdw_ingestion_tools/api/turn_bq/requests/chats.py +++ b/rdw_ingestion_tools/api/turn_bq/requests/chats.py @@ -1,4 +1,4 @@ -from api import concatenate_to_lf +from api import concatenate_to_lazyframe from attrs import define from httpx import Client from polars import LazyFrame @@ -14,7 +14,7 @@ class Chats: client: Client def get_chats_by_id(self, chat_id: int) -> LazyFrame: - """Returns a pandas DataFrame of Turn Chats by chat_id""" + """Returns a Polars LazyFrame of Turn Chats by chat_id""" url = f"chats/{chat_id}" @@ -23,14 +23,14 @@ def get_chats_by_id(self, chat_id: int) -> LazyFrame: url, ) - chats = concatenate_to_lf(chats_generator, chats_schema) + chats = concatenate_to_lazyframe(chats_generator, chats_schema) return chats def get_chats_by_updated_at( self, from_timestamp: str, to_timestamp: str ) -> LazyFrame: - """Returns a pandas DataFrame of Turn Chats by updated_at.""" + """Returns a Polars LazyFrame of Turn Chats by updated_at.""" url = "chats/" params = { @@ -40,6 +40,6 @@ def get_chats_by_updated_at( chats_generator = get_paginated(self.client, url, page_size=1000, **params) - chats = concatenate_to_lf(chats_generator, chats_schema) + chats = concatenate_to_lazyframe(chats_generator, chats_schema) return chats diff --git a/rdw_ingestion_tools/api/turn_bq/requests/contacts.py b/rdw_ingestion_tools/api/turn_bq/requests/contacts.py index e8d16fa..7683fa1 100644 --- a/rdw_ingestion_tools/api/turn_bq/requests/contacts.py +++ b/rdw_ingestion_tools/api/turn_bq/requests/contacts.py @@ -1,4 +1,4 @@ -from api import concatenate_to_lf +from api import concatenate_to_lazyframe from attrs import define from httpx import Client from polars import LazyFrame @@ -14,7 +14,7 @@ class Contacts: client: Client def get_contacts_by_id(self, contact_id: int) -> LazyFrame: - """Returns a pandas DataFrame of Turn Contacts by contact_id""" + """Returns a Polars LazyFrame of Turn Contacts by contact_id""" url = f"contacts/{contact_id}" @@ -23,14 +23,14 @@ def get_contacts_by_id(self, contact_id: int) -> LazyFrame: url, ) - contacts = concatenate_to_lf(contacts_generator, contacts_schema) + contacts = concatenate_to_lazyframe(contacts_generator, contacts_schema) return contacts def get_contacts_by_updated_at( self, from_timestamp: str, to_timestamp: str ) -> LazyFrame: - """Returns a pandas DataFrame of Turn Contacts by updated_at.""" + """Returns a Polars LazyFrame of Turn Contacts by updated_at.""" url = "contacts/" params = { @@ -40,6 +40,6 @@ def get_contacts_by_updated_at( contacts_generator = get_paginated(self.client, url, page_size=1000, **params) - contacts = concatenate_to_lf(contacts_generator, contacts_schema) + contacts = concatenate_to_lazyframe(contacts_generator, contacts_schema) return contacts diff --git a/rdw_ingestion_tools/api/turn_bq/requests/flow_results.py b/rdw_ingestion_tools/api/turn_bq/requests/flow_results.py index 59a3bbb..8c76e45 100644 --- a/rdw_ingestion_tools/api/turn_bq/requests/flow_results.py +++ b/rdw_ingestion_tools/api/turn_bq/requests/flow_results.py @@ -1,4 +1,4 @@ -from api import concatenate_to_lf +from api import concatenate_to_lazyframe from attrs import define from httpx import Client from polars import LazyFrame @@ -19,7 +19,7 @@ class FlowResults: def get_flow_results_by_id(self, stack_uuid: int) -> LazyFrame: """ - Returns a pandas DataFrame of Turn Flow Results + Returns a Polars LazyFrame of Turn Flow Results by stack_uuid """ @@ -31,7 +31,9 @@ def get_flow_results_by_id(self, stack_uuid: int) -> LazyFrame: url, ) - flow_results = concatenate_to_lf(flow_results_generator, flow_results_schema) + flow_results = concatenate_to_lazyframe( + flow_results_generator, flow_results_schema + ) return flow_results @@ -39,7 +41,7 @@ def get_flow_results_by_updated_at( self, from_timestamp: str, to_timestamp: str ) -> LazyFrame: """ - Returns a pandas DataFrame of Turn Flow Results + Returns a Polars LazyFrame of Turn Flow Results by updated_at. """ @@ -55,6 +57,8 @@ def get_flow_results_by_updated_at( self.client, url, page_size=1000, **params ) - flow_results = concatenate_to_lf(flow_results_generator, flow_results_schema) + flow_results = concatenate_to_lazyframe( + flow_results_generator, flow_results_schema + ) return flow_results diff --git a/rdw_ingestion_tools/api/turn_bq/requests/flow_results_data_packages.py b/rdw_ingestion_tools/api/turn_bq/requests/flow_results_data_packages.py index 84f4275..10237aa 100644 --- a/rdw_ingestion_tools/api/turn_bq/requests/flow_results_data_packages.py +++ b/rdw_ingestion_tools/api/turn_bq/requests/flow_results_data_packages.py @@ -1,4 +1,4 @@ -from api import concatenate_to_lf +from api import concatenate_to_lazyframe from attrs import define from httpx import Client from polars import LazyFrame @@ -19,7 +19,7 @@ class FlowResultsDataPackages: def get_flow_results_data_packages_by_id(self, stack_uuid: int) -> LazyFrame: """ - Returns a pandas DataFrame of Turn Flow Results Data Packages + Returns a Polars LazyFrame of Turn Flow Results Data Packages by stack_uuid """ @@ -31,7 +31,7 @@ def get_flow_results_data_packages_by_id(self, stack_uuid: int) -> LazyFrame: url, ) - flow_results_data_packages = concatenate_to_lf( + flow_results_data_packages = concatenate_to_lazyframe( flow_results_data_packages_generator, flow_results_data_packages_schema ) @@ -41,7 +41,7 @@ def get_flow_results_data_packages_by_updated_at( self, from_timestamp: str, to_timestamp: str ) -> LazyFrame: """ - Returns a pandas DataFrame of Turn Flow Results Data Packages + Returns a Polars LazyFrame of Turn Flow Results Data Packages by updated_at. """ @@ -57,7 +57,7 @@ def get_flow_results_data_packages_by_updated_at( self.client, url, page_size=1000, **params ) - flow_results_data_packages = concatenate_to_lf( + flow_results_data_packages = concatenate_to_lazyframe( flow_results_data_packages_generator, flow_results_data_packages_schema ) diff --git a/rdw_ingestion_tools/api/turn_bq/requests/messages.py b/rdw_ingestion_tools/api/turn_bq/requests/messages.py index eb446af..ae2be88 100644 --- a/rdw_ingestion_tools/api/turn_bq/requests/messages.py +++ b/rdw_ingestion_tools/api/turn_bq/requests/messages.py @@ -1,4 +1,4 @@ -from api import concatenate_to_lf +from api import concatenate_to_lazyframe from attrs import define from httpx import Client from polars import LazyFrame @@ -16,7 +16,7 @@ class Messages: def get_messages_by_updated_at( self, from_timestamp: str, to_timestamp: str ) -> LazyFrame: - """Returns a pandas DataFrame of Turn Messages by updated_at.""" + """Returns a Polars LazyFrame of Turn Messages by updated_at.""" url = "messages/" params = { @@ -26,6 +26,6 @@ def get_messages_by_updated_at( messages_generator = get_paginated(self.client, url, page_size=1000, **params) - messages = concatenate_to_lf(messages_generator, messages_schema) + messages = concatenate_to_lazyframe(messages_generator, messages_schema) return messages diff --git a/rdw_ingestion_tools/api/turn_bq/requests/statuses.py b/rdw_ingestion_tools/api/turn_bq/requests/statuses.py index 9e4e757..6bae53c 100644 --- a/rdw_ingestion_tools/api/turn_bq/requests/statuses.py +++ b/rdw_ingestion_tools/api/turn_bq/requests/statuses.py @@ -1,4 +1,4 @@ -from api import concatenate_to_lf +from api import concatenate_to_lazyframe from attrs import define from httpx import Client from polars import LazyFrame @@ -16,7 +16,7 @@ class Statuses: def get_statuses_by_updated_at( self, from_timestamp: str, to_timestamp: str ) -> LazyFrame: - """Returns a pandas DataFrame of Turn Statuses by updated_at.""" + """Returns a Polars LazyFrame of Turn Statuses by updated_at.""" url = "statuses/" params = { @@ -26,6 +26,6 @@ def get_statuses_by_updated_at( status_generator = get_paginated(self.client, url, page_size=1000, **params) - status = concatenate_to_lf(status_generator, statuses_schema) + status = concatenate_to_lazyframe(status_generator, statuses_schema) return status diff --git a/tests/__init__.py b/tests/__init__.py deleted file mode 100644 index e69de29..0000000 diff --git a/tests/conftest.py b/tests/conftest.py index 97a0e1a..77e28cc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,3 +4,5 @@ def pytest_sessionstart(session): os.environ["AAQV2_API_KEY"] = "key" os.environ["AAQV2_API_BASE_URL"] = "http://fake_aaqv2/" + os.environ["RAPIDPRO_API_KEY"] = "RAPIDPRO_API_KEY" + os.environ["RAPIDPRO_API_BASE_URL"] = "http://RAPIDPRO_API_BASE_URL/" diff --git a/tests/fake_aaqv2/api.py b/tests/test_aaqv2/fake_aaqv2/api.py similarity index 100% rename from tests/fake_aaqv2/api.py rename to tests/test_aaqv2/fake_aaqv2/api.py diff --git a/tests/fake_aaqv2/models.py b/tests/test_aaqv2/fake_aaqv2/models.py similarity index 100% rename from tests/fake_aaqv2/models.py rename to tests/test_aaqv2/fake_aaqv2/models.py diff --git a/tests/fake_aaqv2/py.typed b/tests/test_aaqv2/fake_aaqv2/py.typed similarity index 100% rename from tests/fake_aaqv2/py.typed rename to tests/test_aaqv2/fake_aaqv2/py.typed diff --git a/tests/test_aaqv2.py b/tests/test_aaqv2/test_aaqv2.py similarity index 98% rename from tests/test_aaqv2.py rename to tests/test_aaqv2/test_aaqv2.py index 1b36a00..e8d2747 100644 --- a/tests/test_aaqv2.py +++ b/tests/test_aaqv2/test_aaqv2.py @@ -1,18 +1,17 @@ import typing import pytest -from api.aaqv2 import pyAAQV2 from polars import LazyFrame from polars.exceptions import SchemaError from polars.testing import assert_frame_equal +from rdw_ingestion_tools.api.aaqv2 import pyAAQV2 from rdw_ingestion_tools.api.aaqv2.schemas.contents import contents_schema from rdw_ingestion_tools.api.aaqv2.schemas.queries import queries_schema from rdw_ingestion_tools.api.aaqv2.schemas.urgency_queries import urgency_queries_schema from rdw_ingestion_tools.api.aaqv2.schemas.urgency_rules import urgency_rules_schema - -from .fake_aaqv2.api import FakeAAQV2 -from .fake_aaqv2.models import ( +from tests.test_aaqv2.fake_aaqv2.api import FakeAAQV2 +from tests.test_aaqv2.fake_aaqv2.models import ( Content, ContentFeedback, Query, diff --git a/tests/test_concatenation/test_concatenation.py b/tests/test_concatenation/test_concatenation.py new file mode 100644 index 0000000..8ec9306 --- /dev/null +++ b/tests/test_concatenation/test_concatenation.py @@ -0,0 +1,110 @@ +from polars import LazyFrame, Object, String +from polars.testing import assert_frame_equal + +from rdw_ingestion_tools.api import concatenate_to_string_lazyframe, get_polars_schema + + +def test_get_polars_schema_empty_data(): + """ + Tests that schemas generated for empty responses are empty dictionaries. + + """ + schema = get_polars_schema(object_columns=[], data=[]) + + assert schema == {} + + +def test_concatenate_to_string_lazyframe_empty_response(): + """ + Tests that concatenate_to_string_lazyframe returns an empty LazyFrame for + empty response data. + + """ + lf = concatenate_to_string_lazyframe(objs=[], object_columns=[]) + + assert_frame_equal(lf, LazyFrame(schema={})) + + +def test_get_polars_schema_primitive_types(): + """ + Schemas generated from response data use type `String` + for all primitive types. + + """ + data = [ + { + "col1": 1, + "col2": 2.0, + "col3": False, + "col4": None, + "col5": "string", + "col6": "2025-01-01", + } + ] + + schema = get_polars_schema(object_columns=[], data=data) + + expected_schema = { + "col1": String, + "col2": String, + "col3": String, + "col4": String, + "col5": String, + "col6": String, + } + + assert schema == expected_schema + + +def test_get_polars_schema_list_types(): + """ + Tests that generated schemas from response data use type `Object` + for list columns. + + """ + data = [{"col1": [1, 2, 3], "col2": [{"key": "value"}], "col3": False}] + + schema = get_polars_schema(object_columns=["col1", "col2"], data=data) + + expected_schema = {"col1": Object, "col2": Object, "col3": String} + + assert schema == expected_schema + + +def test_get_polars_schema_json_types(): + """ + Tests that generated schemas from response data with JSON columns + normalise the column names in the schema. + + """ + data = [{"col1": {"key": {"inner_key": "value"}}, "col2": {"key": "value"}}] + + schema = get_polars_schema(object_columns=[], data=data) + + expected_schema = {"col1_key_inner_key": String, "col2_key": String} + + assert schema == expected_schema + + +def test_concatenate_to_string_lazyframe(): + """ + Tests that response data is concatenated and normalised into LazyFrames + with column type `String`. + + """ + data = [ + {"col1": 1, "col2": [1, 2, 3], "col3": {"key": "value1"}}, + {"col1": 2, "col2": [1, 2, 3], "col3": {"key": "value2"}}, + ] + + lf = concatenate_to_string_lazyframe(objs=data, object_columns=["col2"]) + + expected_lf = LazyFrame( + { + "col1": ["1", "2"], + "col2": ["[1, 2, 3]", "[1, 2, 3]"], + "col3_key": ["value1", "value2"], + } + ) + + assert_frame_equal(lf, expected_lf) diff --git a/tests/test_pagination/test_rapidpro_pagination.py b/tests/test_pagination/test_rapidpro_pagination.py new file mode 100644 index 0000000..fc8c0af --- /dev/null +++ b/tests/test_pagination/test_rapidpro_pagination.py @@ -0,0 +1,86 @@ +from unittest.mock import Mock + +import pytest + +from rdw_ingestion_tools.api.rapidpro.extensions.httpx import get_paginated + + +@pytest.fixture +def mock_response(): + """Returns a mocked API response object""" + mock_response = Mock() + mock_response.status_code = 200 + mock_response.raise_for_status.return_value = None + + return mock_response + + +@pytest.fixture +def mock_client(mocker): + """Returns a mocked Client object""" + mock_client = mocker.patch( + "rdw_ingestion_tools.api.rapidpro.extensions.httpx.Client" + ) + return mock_client + + +def test_get_paginated_empty_response(mock_client, mock_response): + """ + Tests that the pagination function can handle an empty response. + + """ + mock_response.json.return_value = {"next": None, "previous": None, "results": []} + + mock_client.get.return_value = mock_response + + list(get_paginated(client=mock_client, url="https://test.com")) + + mock_client.get.assert_called_once_with("https://test.com", params={}) + + +def test_get_paginated_cursor_usage(mock_client, mock_response): + """ + Tests that the pagination function call sthe API with the expected cursor. + + """ + + # First mock response object + mock_response.json.return_value = { + "results": [{"id": 1}], + "next": "https://test.com.json?after=2023-01-01+00%3A00%3A00&before=2023-03-01+00%3A00%3A50&cursor=cD0yMDIzLTEwLTI3VDEwJTNBMzAlM0EwMC4wMDAwMDBa&extra_param=1", + } + + # Second mock response object + mock_response2 = Mock() + mock_response2.status_code = 200 + mock_response2.raise_for_status.return_value = None + mock_response2.json.return_value = {"results": [{"id": 3}], "next": None} + + mock_client.get.side_effect = [mock_response, mock_response2] + + generator = get_paginated(client=mock_client, url="https://test.com", foo="bar") + + # First call + next(generator) + mock_client.get.assert_called_with("https://test.com", params={"foo": "bar"}) + + # Second call + next(generator) + mock_client.get.assert_called_with( + "https://test.com", + params={"foo": "bar", "cursor": "cD0yMDIzLTEwLTI3VDEwJTNBMzAlM0EwMC4wMDAwMDBa"}, + ) + + +def test_get_paginated_additional_kwargs(mock_client, mock_response): + """ + Tests that the pagination function can handle additional kwargs. + + """ + mock_response.json.return_value = {"results": [{"id": 3}], "next": None} + + mock_client.get.return_value = mock_response + + list(get_paginated(client=mock_client, url="https://test.com", foo="bar")) + + mock_client.get.assert_called_once_with("https://test.com", params={"foo": "bar"}) diff --git a/uv.lock b/uv.lock index 4ddeec4..8493d80 100644 --- a/uv.lock +++ b/uv.lock @@ -558,6 +558,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5b/4b/d95b052f87db89a2383233c0754c45f6d3b427b7a4bcb771ac9316a6fae1/pytest_coverage-0.0-py2.py3-none-any.whl", hash = "sha256:dedd084c5e74d8e669355325916dc011539b190355021b037242514dee546368", size = 2013, upload-time = "2015-06-17T22:08:36.771Z" }, ] +[[package]] +name = "pytest-mock" +version = "3.15.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/68/14/eb014d26be205d38ad5ad20d9a80f7d201472e08167f0bb4361e251084a9/pytest_mock-3.15.1.tar.gz", hash = "sha256:1849a238f6f396da19762269de72cb1814ab44416fa73a8686deac10b0d87a0f", size = 34036, upload-time = "2025-09-16T16:37:27.081Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5a/cc/06253936f4a7fa2e0f48dfe6d851d9c56df896a9ab09ac019d70b760619c/pytest_mock-3.15.1-py3-none-any.whl", hash = "sha256:0a25e2eb88fe5168d535041d09a4529a188176ae608a6d249ee65abc0949630d", size = 10095, upload-time = "2025-09-16T16:37:25.734Z" }, +] + [[package]] name = "python-dateutil" version = "2.9.0.post0" @@ -639,6 +651,7 @@ dev = [ { name = "mypy" }, { name = "pandas-stubs" }, { name = "pytest-coverage" }, + { name = "pytest-mock" }, { name = "ruff" }, { name = "starlette" }, { name = "tqdm" }, @@ -665,6 +678,7 @@ dev = [ { name = "mypy", specifier = ">=1.10.0" }, { name = "pandas-stubs", specifier = ">=2.2.3.250308" }, { name = "pytest-coverage", specifier = ">=0.0" }, + { name = "pytest-mock", specifier = ">=3.14.1" }, { name = "ruff", specifier = ">=0.4.7" }, { name = "starlette", specifier = ">=0.46.1" }, { name = "tqdm", specifier = ">=4.66.5" },