diff --git a/pipelines/cdp/tests/test_customer_data_platform.py b/pipelines/cdp/tests/test_customer_data_platform.py deleted file mode 100644 index 2cbb6ccf..00000000 --- a/pipelines/cdp/tests/test_customer_data_platform.py +++ /dev/null @@ -1,442 +0,0 @@ -# Copyright 2026 Google LLC -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# https://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Unit tests for Customer Data Platform pipeline transformations and sessionization.""" - -import json -import unittest - -import apache_beam as beam -from apache_beam.options.pipeline_options import GoogleCloudOptions -from apache_beam.testing.test_pipeline import TestPipeline -from apache_beam.testing.util import assert_that -from apache_beam.transforms.window import IntervalWindow -from apache_beam.typehints.schemas import named_tuple_to_schema -from apache_beam.utils.timestamp import Timestamp - -from cdp_pipeline.models import ( - CouponRedemption, - CustomerInteractionEvent, - CustomerSessionProfile, - DeadLetterRecord, - EventType, - TransactionItem, - UnifiedTransactionRecord, -) -from cdp_pipeline.options import MyPipelineOptions -from cdp_pipeline.parsing import ( - ParseRecordDoFn, - TAG_DEADLETTER, -) -from cdp_pipeline.pipeline import build_pipeline -from cdp_pipeline.schemas import load_output_schema -from cdp_pipeline.sessionization import ( - ProcessCustomerSessionDoFn, - TAG_SESSIONS, -) - -# Prevent pytest from treating Apache Beam's TestPipeline as a test case -TestPipeline.__test__ = False - - -class CustomerDataPlatformTest(unittest.TestCase): - - def test_pipeline_options_project(self): - options = MyPipelineOptions(["--project=my-test-project"]) - gcp_options = options.view_as(GoogleCloudOptions) - self.assertEqual(gcp_options.project, "my-test-project") - - def test_load_output_schema_default(self): - schema = load_output_schema() - self.assertIn("fields", schema) - field_names = [field["name"] for field in schema["fields"]] - self.assertIn("transaction_id", field_names) - self.assertIn("household_key", field_names) - self.assertIn("coupon_upc", field_names) - self.assertIn("product_id", field_names) - self.assertIn("coupon_discount", field_names) - self.assertIn("session_id", field_names) - self.assertIn("campaign", field_names) - - def test_load_schemas_helpers(self): - sessions_schema = load_output_schema(None, "customer_sessions.json") - self.assertIn("fields", sessions_schema) - session_fields = [f["name"] for f in sessions_schema["fields"]] - self.assertIn("session_id", session_fields) - self.assertIn("total_spend", session_fields) - self.assertIn("total_transactions", session_fields) - - dlq_schema = load_output_schema(None, "deadletter_table.json") - self.assertIn("fields", dlq_schema) - dlq_fields = [f["name"] for f in dlq_schema["fields"]] - self.assertIn("error_message", dlq_fields) - self.assertIn("raw_payload", dlq_fields) - - with self.assertRaises(FileNotFoundError): - load_output_schema(None, "non_existent_schema.json") - - def test_parse_record_valid_transaction(self): - fn = ParseRecordDoFn(EventType.TRANSACTION) - fn.setup() - raw = json.dumps({ - "household_key": "100", - "transaction_id": "tx-1", - "product_id": "prod-1", - "sales_value": 15.50 - }).encode("utf-8") - - results = list(fn.process(raw)) - self.assertEqual(len(results), 1) - hh_key, event = results[0] - self.assertEqual(hh_key, "100") - self.assertIsInstance(event, CustomerInteractionEvent) - self.assertEqual(event.transaction_id, "tx-1") - self.assertEqual(event.event_type, EventType.TRANSACTION.value) - self.assertIsNotNone(event.transaction) - self.assertEqual(event.transaction.product_id, "prod-1") - self.assertEqual(event.transaction.sales_value, 15.50) - self.assertIsNone(event.coupon) - - def test_parse_record_valid_coupon(self): - fn = ParseRecordDoFn(EventType.COUPON) - fn.setup() - raw = json.dumps({ - "household_key": "100", - "transaction_id": "tx-1", - "coupon_upc": "cp-1", - "campaign": "camp-99" - }) - - results = list(fn.process(raw)) - self.assertEqual(len(results), 1) - hh_key, event = results[0] - self.assertEqual(hh_key, "100") - self.assertIsInstance(event, CustomerInteractionEvent) - self.assertEqual(event.transaction_id, "tx-1") - self.assertEqual(event.event_type, EventType.COUPON.value) - self.assertIsNotNone(event.coupon) - self.assertEqual(event.coupon.coupon_upc, "cp-1") - self.assertEqual(event.coupon.campaign, "camp-99") - self.assertIsNone(event.transaction) - - def test_parse_record_malformed_json_dlq(self): - fn = ParseRecordDoFn("transaction") - fn.setup() - bad_bytes = b"BROKEN_JSON_DATA{{{" - - results = list(fn.process(bad_bytes)) - self.assertEqual(len(results), 1) - tagged_output = results[0] - self.assertIsInstance(tagged_output, beam.pvalue.TaggedOutput) - self.assertEqual(tagged_output.tag, TAG_DEADLETTER) - self.assertIn("Malformed payload", tagged_output.value["error_message"]) - - def test_parse_record_missing_keys_dlq(self): - fn = ParseRecordDoFn("transaction") - fn.setup() - # Missing household_key - missing_key = json.dumps({"transaction_id": "tx-1"}).encode("utf-8") - - results = list(fn.process(missing_key)) - self.assertEqual(len(results), 1) - tagged_output = results[0] - self.assertIsInstance(tagged_output, beam.pvalue.TaggedOutput) - self.assertEqual(tagged_output.tag, TAG_DEADLETTER) - self.assertIn("Missing required household_key", - tagged_output.value["error_message"]) - - def test_process_customer_session_aggregation(self): - fn = ProcessCustomerSessionDoFn() - fn.setup() - - household_key = "hh-42" - mock_window = IntervalWindow( # pylint: disable=too-many-function-args - Timestamp(1000), Timestamp(1300)) - - events = [ - CustomerInteractionEvent( - event_type=EventType.TRANSACTION.value, - household_key="hh-42", - transaction_id="tx-101", - transaction=TransactionItem( - product_id="prod-A", - quantity=2, - sales_value=20.0, - retail_disc=2.0, - coupon_disc=1.0, - store_id="store-1", - ), - ), - CustomerInteractionEvent( - event_type=EventType.TRANSACTION.value, - household_key="hh-42", - transaction_id="tx-102", - transaction=TransactionItem( - product_id="prod-B", - quantity=1, - sales_value=10.0, - retail_disc=0.0, - coupon_disc=0.0, - store_id="store-1", - ), - ), - CustomerInteractionEvent( - event_type=EventType.COUPON.value, - household_key="hh-42", - transaction_id="tx-101", - coupon=CouponRedemption( - coupon_upc="cp-999", - campaign="fall-sale", - ), - ), - ] - - outputs = list(fn.process((household_key, events), window=mock_window)) - - # Main outputs: unified transaction records - unified = [ - item for item in outputs - if not isinstance(item, beam.pvalue.TaggedOutput) - ] - # Tagged output: Customer 360 session profile - sessions = [ - item.value for item in outputs if - isinstance(item, beam.pvalue.TaggedOutput) and item.tag == TAG_SESSIONS - ] - - # Two transaction items were emitted - self.assertEqual(len(unified), 2) - self.assertIsInstance(unified[0], UnifiedTransactionRecord) - tx_101 = next(u for u in unified if u.transaction_id == "tx-101") - self.assertEqual(tx_101.coupon_upc, "cp-999") - self.assertEqual(tx_101.campaign, "fall-sale") - self.assertEqual(tx_101.sales_value, 20.0) - - tx_102 = next(u for u in unified if u.transaction_id == "tx-102") - self.assertIsNone(tx_102.coupon_upc) - - # Verify session summary - self.assertEqual(len(sessions), 1) - session = sessions[0] - self.assertIsInstance(session, CustomerSessionProfile) - self.assertEqual(session.household_key, "hh-42") - self.assertEqual(session.total_transactions, 2) - self.assertEqual(session.total_items_purchased, 3) - self.assertEqual(session.total_spend, 30.0) - self.assertEqual(session.total_discount, 3.0) - self.assertEqual(session.coupons_redeemed_count, 1) - self.assertEqual(session.distinct_products_count, 2) - self.assertIn("fall-sale", session.campaigns) - - def test_build_pipeline_end_to_end_in_memory(self): - options = MyPipelineOptions( - session_gap_seconds=10, - allowed_lateness_seconds=5, - output_dataset="test_dataset", - output_table="test_unified", - output_sessions_table="test_sessions", - ) - - in_memory_tx = [ - json.dumps({ - "household_key": "hh-1", - "transaction_id": "tx-1", - "product_id": "p100", - "quantity": 1, - "sales_value": 5.0, - "event_timestamp": "2026-09-08T10:00:00Z", - }).encode("utf-8"), - b"MALFORMED_JSON_PAYLOAD", - ] - in_memory_cp = [ - json.dumps({ - "household_key": "hh-1", - "transaction_id": "tx-1", - "coupon_upc": "c100", - "campaign": "summer", - "event_timestamp": "2026-09-08T10:00:02Z", - }).encode("utf-8") - ] - - with TestPipeline(options=options) as p: - unified, sessions, deadletters = build_pipeline( - pipeline=p, - pipeline_options=options, - in_memory_transactions=in_memory_tx, - in_memory_coupons=in_memory_cp, - ) - - assert_that(unified, _check_unified, label="CheckUnified") - assert_that(sessions, _check_sessions, label="CheckSessions") - assert_that(deadletters, _check_deadletters, label="CheckDeadletters") - - -def _check_unified(records): - assert len(records) == 1 - assert records[0].transaction_id == "tx-1" - assert records[0].coupon_upc == "c100" - assert records[0].campaign == "summer" - - -def _check_sessions(session_records): - assert len(session_records) == 1 - assert session_records[0].household_key == "hh-1" - assert session_records[0].total_spend == 5.0 - - -def _check_deadletters(dlq_records): - assert len(dlq_records) == 1 - assert dlq_records[0]["source"] == "transaction" - assert "Malformed payload" in dlq_records[0]["error_message"] - - -class ModelsTest(unittest.TestCase): - """Unit tests for Beam schema data models and parsing logic.""" - - def test_transaction_item_defaults_and_dict(self): - item = TransactionItem( - product_id="prod-1", - quantity=3, - sales_value=25.50, - store_id="store-10", - retail_disc=1.50, - coupon_disc=0.50, - ) - self.assertEqual(item.product_id, "prod-1") - self.assertEqual(item.quantity, 3) - self.assertEqual(item.sales_value, 25.50) - item_dict = item.to_dict() - self.assertEqual(item_dict["product_id"], "prod-1") - self.assertEqual(item_dict["quantity"], 3) - self.assertEqual(item_dict["sales_value"], 25.50) - self.assertEqual(item_dict["store_id"], "store-10") - self.assertIsNone(item_dict["day"]) - - def test_coupon_redemption_defaults_and_dict(self): - coupon = CouponRedemption( - coupon_upc="cp-12345", - campaign="spring-sale", - day=15, - ) - self.assertEqual(coupon.coupon_upc, "cp-12345") - self.assertEqual(coupon.campaign, "spring-sale") - self.assertEqual(coupon.day, 15) - c_dict = coupon.to_dict() - self.assertEqual(c_dict["coupon_upc"], "cp-12345") - self.assertEqual(c_dict["campaign"], "spring-sale") - self.assertEqual(c_dict["day"], 15) - - def test_dead_letter_record_dict(self): - dlq = DeadLetterRecord( - source="transaction", - raw_payload='{"bad": "data"}', - error_message="Missing required field", - timestamp="2026-09-07T12:00:00Z", - ) - self.assertEqual(dlq.source, "transaction") - d_dict = dlq.to_dict() - self.assertEqual(d_dict["source"], "transaction") - self.assertEqual(d_dict["error_message"], "Missing required field") - - def test_customer_interaction_event_from_raw_payload_transaction(self): - payload = json.dumps({ - "household_key": "hh-99", - "transaction_id": "tx-888", - "product_id": "prod-abc", - "quantity": "2", - "sales_value": "19.99", - "retail_disc": "1.00", - "coupon_disc": "0.50", - "store_id": "st-5", - }).encode("utf-8") - - event, dlq = CustomerInteractionEvent.from_raw_payload( - payload, EventType.TRANSACTION) - self.assertIsNone(dlq) - self.assertIsNotNone(event) - self.assertEqual(event.household_key, "hh-99") - self.assertEqual(event.transaction_id, "tx-888") - self.assertEqual(event.event_type, EventType.TRANSACTION.value) - self.assertIsNotNone(event.transaction) - self.assertEqual(event.transaction.product_id, "prod-abc") - self.assertEqual(event.transaction.quantity, 2) - self.assertEqual(event.transaction.sales_value, 19.99) - self.assertEqual(event.transaction.retail_disc, 1.00) - self.assertEqual(event.transaction.coupon_disc, 0.50) - self.assertIsNone(event.coupon) - - event_dict = event.to_dict() - self.assertEqual(event_dict["household_key"], "hh-99") - self.assertIsNotNone(event_dict["transaction"]) - self.assertIsNone(event_dict["coupon"]) - - def test_customer_interaction_event_from_raw_payload_coupon(self): - payload = { - "household_key": "hh-99", - "transaction_id": "tx-888", - "coupon_upc": "cp-99999", - "campaign": "promo-2026", - "day": 42, - } - - event, dlq = CustomerInteractionEvent.from_raw_payload( - payload, EventType.COUPON) - self.assertIsNone(dlq) - self.assertIsNotNone(event) - self.assertEqual(event.household_key, "hh-99") - self.assertEqual(event.transaction_id, "tx-888") - self.assertEqual(event.event_type, EventType.COUPON.value) - self.assertIsNotNone(event.coupon) - self.assertEqual(event.coupon.coupon_upc, "cp-99999") - self.assertEqual(event.coupon.campaign, "promo-2026") - self.assertEqual(event.coupon.day, 42) - self.assertIsNone(event.transaction) - - def test_customer_interaction_event_from_raw_payload_malformed_json(self): - event, dlq = CustomerInteractionEvent.from_raw_payload( - b"NOT_A_JSON_STRING", EventType.TRANSACTION) - self.assertIsNone(event) - self.assertIsNotNone(dlq) - self.assertEqual(dlq.source, EventType.TRANSACTION.value) - self.assertIn("Malformed payload", dlq.error_message) - - def test_customer_interaction_event_from_raw_payload_missing_keys(self): - payload = json.dumps({"product_id": "prod-1"}) - event, dlq = CustomerInteractionEvent.from_raw_payload( - payload, EventType.TRANSACTION) - self.assertIsNone(event) - self.assertIsNotNone(dlq) - self.assertIn("Missing required household_key", dlq.error_message) - - def test_customer_interaction_event_from_raw_payload_unsupported_type(self): - event, dlq = CustomerInteractionEvent.from_raw_payload( - 12345, EventType.TRANSACTION) # type: ignore[arg-type] - self.assertIsNone(event) - self.assertIsNotNone(dlq) - self.assertIn("Unsupported payload type", dlq.error_message) - - def test_beam_schema_compatibility(self): - for model_cls in ( - TransactionItem, - CouponRedemption, - CustomerInteractionEvent, - UnifiedTransactionRecord, - CustomerSessionProfile, - ): - schema = named_tuple_to_schema(model_cls) - self.assertIsNotNone(schema) - self.assertGreater(len(schema.fields), 0) - - -if __name__ == "__main__": - unittest.main() diff --git a/pipelines/cdp/tests/test_models.py b/pipelines/cdp/tests/test_models.py new file mode 100644 index 00000000..e0eefaa8 --- /dev/null +++ b/pipelines/cdp/tests/test_models.py @@ -0,0 +1,171 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for Beam schema data models and event payload parsing.""" + +import json +import unittest + +from apache_beam.typehints.schemas import named_tuple_to_schema + +from cdp_pipeline.models import ( + CouponRedemption, + CustomerInteractionEvent, + CustomerSessionProfile, + DeadLetterRecord, + EventType, + TransactionItem, + UnifiedTransactionRecord, +) + + +class ModelsTest(unittest.TestCase): + """Unit tests for Beam schema data models and parsing logic.""" + + def test_transaction_item_defaults_and_dict(self): + item = TransactionItem( + product_id="prod-1", + quantity=3, + sales_value=25.50, + store_id="store-10", + retail_disc=1.50, + coupon_disc=0.50, + ) + self.assertEqual(item.product_id, "prod-1") + self.assertEqual(item.quantity, 3) + self.assertEqual(item.sales_value, 25.50) + item_dict = item.to_dict() + self.assertEqual(item_dict["product_id"], "prod-1") + self.assertEqual(item_dict["quantity"], 3) + self.assertEqual(item_dict["sales_value"], 25.50) + self.assertEqual(item_dict["store_id"], "store-10") + self.assertIsNone(item_dict["day"]) + + def test_coupon_redemption_defaults_and_dict(self): + coupon = CouponRedemption( + coupon_upc="cp-12345", + campaign="spring-sale", + day=15, + ) + self.assertEqual(coupon.coupon_upc, "cp-12345") + self.assertEqual(coupon.campaign, "spring-sale") + self.assertEqual(coupon.day, 15) + c_dict = coupon.to_dict() + self.assertEqual(c_dict["coupon_upc"], "cp-12345") + self.assertEqual(c_dict["campaign"], "spring-sale") + self.assertEqual(c_dict["day"], 15) + + def test_dead_letter_record_dict(self): + dlq = DeadLetterRecord( + source="transaction", + raw_payload='{"bad": "data"}', + error_message="Missing required field", + timestamp="2026-09-07T12:00:00Z", + ) + self.assertEqual(dlq.source, "transaction") + d_dict = dlq.to_dict() + self.assertEqual(d_dict["source"], "transaction") + self.assertEqual(d_dict["error_message"], "Missing required field") + + def test_customer_interaction_event_from_raw_payload_transaction(self): + payload = json.dumps({ + "household_key": "hh-99", + "transaction_id": "tx-888", + "product_id": "prod-abc", + "quantity": "2", + "sales_value": "19.99", + "retail_disc": "1.00", + "coupon_disc": "0.50", + "store_id": "st-5", + }).encode("utf-8") + + event, dlq = CustomerInteractionEvent.from_raw_payload( + payload, EventType.TRANSACTION) + self.assertIsNone(dlq) + self.assertIsNotNone(event) + self.assertEqual(event.household_key, "hh-99") + self.assertEqual(event.transaction_id, "tx-888") + self.assertEqual(event.event_type, EventType.TRANSACTION.value) + self.assertIsNotNone(event.transaction) + self.assertEqual(event.transaction.product_id, "prod-abc") + self.assertEqual(event.transaction.quantity, 2) + self.assertEqual(event.transaction.sales_value, 19.99) + self.assertEqual(event.transaction.retail_disc, 1.00) + self.assertEqual(event.transaction.coupon_disc, 0.50) + self.assertIsNone(event.coupon) + + event_dict = event.to_dict() + self.assertEqual(event_dict["household_key"], "hh-99") + self.assertIsNotNone(event_dict["transaction"]) + self.assertIsNone(event_dict["coupon"]) + + def test_customer_interaction_event_from_raw_payload_coupon(self): + payload = { + "household_key": "hh-99", + "transaction_id": "tx-888", + "coupon_upc": "cp-99999", + "campaign": "promo-2026", + "day": 42, + } + + event, dlq = CustomerInteractionEvent.from_raw_payload( + payload, EventType.COUPON) + self.assertIsNone(dlq) + self.assertIsNotNone(event) + self.assertEqual(event.household_key, "hh-99") + self.assertEqual(event.transaction_id, "tx-888") + self.assertEqual(event.event_type, EventType.COUPON.value) + self.assertIsNotNone(event.coupon) + self.assertEqual(event.coupon.coupon_upc, "cp-99999") + self.assertEqual(event.coupon.campaign, "promo-2026") + self.assertEqual(event.coupon.day, 42) + self.assertIsNone(event.transaction) + + def test_customer_interaction_event_from_raw_payload_malformed_json(self): + event, dlq = CustomerInteractionEvent.from_raw_payload( + b"NOT_A_JSON_STRING", EventType.TRANSACTION) + self.assertIsNone(event) + self.assertIsNotNone(dlq) + self.assertEqual(dlq.source, EventType.TRANSACTION.value) + self.assertIn("Malformed payload", dlq.error_message) + + def test_customer_interaction_event_from_raw_payload_missing_keys(self): + payload = json.dumps({"product_id": "prod-1"}) + event, dlq = CustomerInteractionEvent.from_raw_payload( + payload, EventType.TRANSACTION) + self.assertIsNone(event) + self.assertIsNotNone(dlq) + self.assertIn("Missing required household_key", dlq.error_message) + + def test_customer_interaction_event_from_raw_payload_unsupported_type(self): + event, dlq = CustomerInteractionEvent.from_raw_payload( + 12345, EventType.TRANSACTION) # type: ignore[arg-type] + self.assertIsNone(event) + self.assertIsNotNone(dlq) + self.assertIn("Unsupported payload type", dlq.error_message) + + def test_beam_schema_compatibility(self): + for model_cls in ( + TransactionItem, + CouponRedemption, + CustomerInteractionEvent, + UnifiedTransactionRecord, + CustomerSessionProfile, + ): + schema = named_tuple_to_schema(model_cls) + self.assertIsNotNone(schema) + self.assertGreater(len(schema.fields), 0) + + +if __name__ == "__main__": + unittest.main() diff --git a/pipelines/cdp/tests/test_options.py b/pipelines/cdp/tests/test_options.py new file mode 100644 index 00000000..3e19b7a8 --- /dev/null +++ b/pipelines/cdp/tests/test_options.py @@ -0,0 +1,84 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for Customer Data Platform pipeline options.""" + +import unittest + +from apache_beam.options.pipeline_options import GoogleCloudOptions + +from cdp_pipeline.options import MyPipelineOptions + + +class OptionsTest(unittest.TestCase): + """Unit tests for MyPipelineOptions CLI argument parsing and defaults.""" + + def test_pipeline_options_defaults(self): + options = MyPipelineOptions([]) + self.assertIsNone(options.transactions_topic) + self.assertIsNone(options.transactions_subscription) + self.assertIsNone(options.coupons_redemption_topic) + self.assertIsNone(options.coupons_redemption_subscription) + self.assertEqual(options.output_dataset, "cdp_dataset") + self.assertEqual(options.output_table, "unified_customer_data") + self.assertEqual(options.output_sessions_table, "customer_sessions") + self.assertIsNone(options.deadletter_table) + self.assertEqual(options.session_gap_seconds, 900) + self.assertEqual(options.allowed_lateness_seconds, 60) + self.assertTrue(options.use_storage_write_api) + self.assertIsNone(options.output_schema_path) + self.assertIsNone(options.output_sessions_schema_path) + self.assertIsNone(options.deadletter_schema_path) + + def test_pipeline_options_project(self): + options = MyPipelineOptions(["--project=my-test-project"]) + gcp_options = options.view_as(GoogleCloudOptions) + self.assertEqual(gcp_options.project, "my-test-project") + + def test_pipeline_options_custom_flags(self): + flags = [ + "--transactions_topic=projects/p/topics/tx", + "--transactions_subscription=projects/p/subscriptions/tx-sub", + "--coupons_redemption_topic=projects/p/topics/cp", + "--coupons_redemption_subscription=projects/p/subscriptions/cp-sub", + "--output_dataset=custom_dataset", + "--output_table=custom_unified", + "--output_sessions_table=custom_sessions", + "--deadletter_table=custom_dlq", + "--session_gap_seconds=300", + "--allowed_lateness_seconds=30", + "--output_schema_path=/path/to/unified.json", + "--output_sessions_schema_path=/path/to/sessions.json", + "--deadletter_schema_path=/path/to/dlq.json", + ] + options = MyPipelineOptions(flags) + self.assertEqual(options.transactions_topic, "projects/p/topics/tx") + self.assertEqual(options.transactions_subscription, + "projects/p/subscriptions/tx-sub") + self.assertEqual(options.coupons_redemption_topic, "projects/p/topics/cp") + self.assertEqual(options.coupons_redemption_subscription, + "projects/p/subscriptions/cp-sub") + self.assertEqual(options.output_dataset, "custom_dataset") + self.assertEqual(options.output_table, "custom_unified") + self.assertEqual(options.output_sessions_table, "custom_sessions") + self.assertEqual(options.deadletter_table, "custom_dlq") + self.assertEqual(options.session_gap_seconds, 300) + self.assertEqual(options.allowed_lateness_seconds, 30) + self.assertEqual(options.output_schema_path, "/path/to/unified.json") + self.assertEqual(options.output_sessions_schema_path, + "/path/to/sessions.json") + self.assertEqual(options.deadletter_schema_path, "/path/to/dlq.json") + + +if __name__ == "__main__": + unittest.main() diff --git a/pipelines/cdp/tests/test_parsing.py b/pipelines/cdp/tests/test_parsing.py new file mode 100644 index 00000000..2a862a8f --- /dev/null +++ b/pipelines/cdp/tests/test_parsing.py @@ -0,0 +1,183 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for record parsing, validation, and timestamp assignment transforms.""" + +from datetime import datetime, timezone +import json +import unittest + +import apache_beam as beam +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that, equal_to + +from cdp_pipeline.models import ( + CustomerInteractionEvent, + EventType, +) +from cdp_pipeline.parsing import ( + AssignEventTimestampDoFn, + ParseRecordDoFn, + TAG_DEADLETTER, +) + +# Prevent pytest from treating Apache Beam's TestPipeline as a test case +TestPipeline.__test__ = False + + +class ParsingTest(unittest.TestCase): + """Unit tests for ParseRecordDoFn and AssignEventTimestampDoFn.""" + + def test_parse_record_valid_transaction(self): + raw = json.dumps({ + "household_key": "100", + "transaction_id": "tx-1", + "product_id": "prod-1", + "sales_value": 15.50 + }).encode("utf-8") + + with TestPipeline() as p: + results = ( + p + | beam.Create([raw]) + | beam.ParDo(ParseRecordDoFn(EventType.TRANSACTION)).with_outputs( + TAG_DEADLETTER, main="valid")) + + def check_valid(elements): + assert len(elements) == 1 + hh_key, event = elements[0] + assert hh_key == "100" + assert event.transaction_id == "tx-1" + assert event.event_type == EventType.TRANSACTION.value + assert event.transaction is not None + assert event.transaction.product_id == "prod-1" + assert event.transaction.sales_value == 15.50 + assert event.coupon is None + + assert_that(results.valid, check_valid) + assert_that(results[TAG_DEADLETTER], equal_to([])) + + def test_parse_record_valid_coupon(self): + raw = json.dumps({ + "household_key": "100", + "transaction_id": "tx-1", + "coupon_upc": "cp-1", + "campaign": "camp-99" + }) + + with TestPipeline() as p: + results = ( + p + | beam.Create([raw]) + | beam.ParDo(ParseRecordDoFn(EventType.COUPON)).with_outputs( + TAG_DEADLETTER, main="valid")) + + def check_valid(elements): + assert len(elements) == 1 + hh_key, event = elements[0] + assert hh_key == "100" + assert event.transaction_id == "tx-1" + assert event.event_type == EventType.COUPON.value + assert event.coupon is not None + assert event.coupon.coupon_upc == "cp-1" + assert event.coupon.campaign == "camp-99" + assert event.transaction is None + + assert_that(results.valid, check_valid) + assert_that(results[TAG_DEADLETTER], equal_to([])) + + def test_parse_record_string_event_type(self): + fn = ParseRecordDoFn("transaction") + self.assertEqual(fn.record_type, EventType.TRANSACTION) + + def test_parse_record_malformed_json_dlq(self): + bad_bytes = b"BROKEN_JSON_DATA{{{" + + with TestPipeline() as p: + results = ( + p + | beam.Create([bad_bytes]) + | beam.ParDo(ParseRecordDoFn("transaction")).with_outputs( + TAG_DEADLETTER, main="valid")) + + def check_dlq(elements): + assert len(elements) == 1 + assert "Malformed payload" in elements[0]["error_message"] + + assert_that(results.valid, equal_to([])) + assert_that(results[TAG_DEADLETTER], check_dlq) + + def test_parse_record_missing_keys_dlq(self): + missing_key = json.dumps({"transaction_id": "tx-1"}).encode("utf-8") + + with TestPipeline() as p: + results = ( + p + | beam.Create([missing_key]) + | beam.ParDo(ParseRecordDoFn("transaction")).with_outputs( + TAG_DEADLETTER, main="valid")) + + def check_dlq(elements): + assert len(elements) == 1 + assert "Missing required household_key" in elements[0]["error_message"] + + assert_that(results.valid, equal_to([])) + assert_that(results[TAG_DEADLETTER], check_dlq) + + def test_assign_event_timestamp_with_iso_string(self): + iso_ts = "2026-09-08T10:00:00Z" + event = CustomerInteractionEvent( + event_type=EventType.TRANSACTION.value, + household_key="hh-1", + transaction_id="tx-1", + event_timestamp=iso_ts, + ) + expected_seconds = datetime.fromisoformat( + "2026-09-08T10:00:00+00:00").timestamp() + + with TestPipeline() as p: + output = ( + p + | beam.Create([("hh-1", event)]) + | beam.ParDo(AssignEventTimestampDoFn()) + | beam.Map(lambda el, ts=beam.DoFn.TimestampParam: float(ts.micros) / + 1000000.0)) + + assert_that(output, equal_to([expected_seconds])) + + def test_assign_event_timestamp_fallback_to_utc_now(self): + event = CustomerInteractionEvent( + event_type=EventType.TRANSACTION.value, + household_key="hh-1", + transaction_id="tx-1", + event_timestamp="INVALID_DATE_STRING", + ) + before_ts = datetime.now(timezone.utc).timestamp() + + with TestPipeline() as p: + output = ( + p + | beam.Create([("hh-1", event)]) + | beam.ParDo(AssignEventTimestampDoFn()) + | beam.Map(lambda el, ts=beam.DoFn.TimestampParam: float(ts.micros) / + 1000000.0)) + + def check_fallback(elements): + assert len(elements) == 1 + assert elements[0] >= before_ts + + assert_that(output, check_fallback) + + +if __name__ == "__main__": + unittest.main() diff --git a/pipelines/cdp/tests/test_pipeline.py b/pipelines/cdp/tests/test_pipeline.py new file mode 100644 index 00000000..9a52675b --- /dev/null +++ b/pipelines/cdp/tests/test_pipeline.py @@ -0,0 +1,139 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Integration and DAG construction tests for the Customer Data Platform pipeline.""" + +import json +import unittest + +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that, equal_to + +from cdp_pipeline.options import MyPipelineOptions +from cdp_pipeline.pipeline import build_pipeline + +# Prevent pytest from treating Apache Beam's TestPipeline as a test case +TestPipeline.__test__ = False + + +def _check_unified(records): + assert len(records) == 1 + assert records[0].transaction_id == "tx-1" + assert records[0].coupon_upc == "c100" + assert records[0].campaign == "summer" + + +def _check_sessions(session_records): + assert len(session_records) == 1 + assert session_records[0].household_key == "hh-1" + assert session_records[0].total_spend == 5.0 + + +def _check_deadletters(dlq_records): + assert len(dlq_records) == 1 + assert dlq_records[0]["source"] == "transaction" + assert "Malformed payload" in dlq_records[0]["error_message"] + + +class PipelineTest(unittest.TestCase): + """Integration and DAG assembly tests for the CDP streaming pipeline.""" + + def test_build_pipeline_end_to_end_in_memory(self): + options = MyPipelineOptions( + session_gap_seconds=10, + allowed_lateness_seconds=5, + output_dataset="test_dataset", + output_table="test_unified", + output_sessions_table="test_sessions", + ) + + in_memory_tx = [ + json.dumps({ + "household_key": "hh-1", + "transaction_id": "tx-1", + "product_id": "p100", + "quantity": 1, + "sales_value": 5.0, + "event_timestamp": "2026-09-08T10:00:00Z", + }).encode("utf-8"), + b"MALFORMED_JSON_PAYLOAD", + ] + in_memory_cp = [ + json.dumps({ + "household_key": "hh-1", + "transaction_id": "tx-1", + "coupon_upc": "c100", + "campaign": "summer", + "event_timestamp": "2026-09-08T10:00:02Z", + }).encode("utf-8") + ] + + with TestPipeline(options=options) as p: + unified, sessions, deadletters = build_pipeline( + pipeline=p, + pipeline_options=options, + in_memory_transactions=in_memory_tx, + in_memory_coupons=in_memory_cp, + ) + + assert_that(unified, _check_unified, label="CheckUnified") + assert_that(sessions, _check_sessions, label="CheckSessions") + assert_that(deadletters, _check_deadletters, label="CheckDeadletters") + + def test_build_pipeline_empty_inputs(self): + options = MyPipelineOptions( + session_gap_seconds=10, + allowed_lateness_seconds=0, + ) + with TestPipeline(options=options) as p: + unified, sessions, deadletters = build_pipeline( + pipeline=p, + pipeline_options=options, + in_memory_transactions=[], + in_memory_coupons=[], + ) + assert_that(unified, equal_to([]), label="CheckEmptyUnified") + assert_that(sessions, equal_to([]), label="CheckEmptySessions") + assert_that(deadletters, equal_to([]), label="CheckEmptyDeadletters") + + def test_build_pipeline_pubsub_topic_branch_dag(self): + options = MyPipelineOptions([ + "--transactions_topic=projects/test-p/topics/tx-topic", + "--coupons_redemption_topic=projects/test-p/topics/cp-topic", + ]) + p = TestPipeline(options=options) + unified, sessions, deadletters = build_pipeline( + pipeline=p, + pipeline_options=options, + ) + self.assertIsNotNone(unified) + self.assertIsNotNone(sessions) + self.assertIsNotNone(deadletters) + + def test_build_pipeline_pubsub_subscription_branch_dag(self): + options = MyPipelineOptions([ + "--transactions_subscription=projects/test-p/subscriptions/tx-sub", + "--coupons_redemption_subscription=projects/test-p/subscriptions/cp-sub", + ]) + p = TestPipeline(options=options) + unified, sessions, deadletters = build_pipeline( + pipeline=p, + pipeline_options=options, + ) + self.assertIsNotNone(unified) + self.assertIsNotNone(sessions) + self.assertIsNotNone(deadletters) + + +if __name__ == "__main__": + unittest.main() diff --git a/pipelines/cdp/tests/test_schemas.py b/pipelines/cdp/tests/test_schemas.py new file mode 100644 index 00000000..b925e83b --- /dev/null +++ b/pipelines/cdp/tests/test_schemas.py @@ -0,0 +1,81 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for BigQuery table schema loading.""" + +import json +import tempfile +import unittest + +from cdp_pipeline.schemas import load_output_schema + + +class SchemasTest(unittest.TestCase): + """Unit tests for load_output_schema with packaged and custom schemas.""" + + def test_load_output_schema_default(self): + schema = load_output_schema() + self.assertIn("fields", schema) + field_names = [field["name"] for field in schema["fields"]] + self.assertIn("transaction_id", field_names) + self.assertIn("household_key", field_names) + self.assertIn("coupon_upc", field_names) + self.assertIn("product_id", field_names) + self.assertIn("coupon_discount", field_names) + self.assertIn("session_id", field_names) + self.assertIn("campaign", field_names) + self.assertIn("sales_value", field_names) + self.assertIn("quantity", field_names) + self.assertIn("processed_timestamp", field_names) + + def test_load_schemas_helpers(self): + sessions_schema = load_output_schema(None, "customer_sessions.json") + self.assertIn("fields", sessions_schema) + session_fields = [f["name"] for f in sessions_schema["fields"]] + self.assertIn("session_id", session_fields) + self.assertIn("household_key", session_fields) + self.assertIn("total_spend", session_fields) + self.assertIn("total_transactions", session_fields) + self.assertIn("total_items_purchased", session_fields) + self.assertIn("coupons_redeemed_count", session_fields) + + dlq_schema = load_output_schema(None, "deadletter_table.json") + self.assertIn("fields", dlq_schema) + dlq_fields = [f["name"] for f in dlq_schema["fields"]] + self.assertIn("source", dlq_fields) + self.assertIn("error_message", dlq_fields) + self.assertIn("raw_payload", dlq_fields) + self.assertIn("timestamp", dlq_fields) + + def test_load_output_schema_non_existent(self): + with self.assertRaises(FileNotFoundError): + load_output_schema(None, "non_existent_schema.json") + + def test_load_output_schema_custom_path(self): + custom_content = { + "fields": [{ + "name": "custom_field", + "type": "STRING", + "mode": "REQUIRED" + }] + } + with tempfile.NamedTemporaryFile( + "w+", encoding="utf-8", delete=True) as temp_file: + json.dump(custom_content, temp_file) + temp_file.flush() + loaded = load_output_schema(temp_file.name) + self.assertEqual(loaded, custom_content) + + +if __name__ == "__main__": + unittest.main() diff --git a/pipelines/cdp/tests/test_sessionization.py b/pipelines/cdp/tests/test_sessionization.py new file mode 100644 index 00000000..b9dbf336 --- /dev/null +++ b/pipelines/cdp/tests/test_sessionization.py @@ -0,0 +1,233 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for customer session aggregation and session profile generation.""" + +import unittest + +import apache_beam as beam +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.testing.util import assert_that, equal_to +from apache_beam.transforms.window import Sessions, TimestampedValue + +from cdp_pipeline.models import ( + CouponRedemption, + CustomerInteractionEvent, + CustomerSessionProfile, + EventType, + TransactionItem, + UnifiedTransactionRecord, +) +from cdp_pipeline.sessionization import ( + ProcessCustomerSessionDoFn, + TAG_SESSIONS, +) + +# Prevent pytest from treating Apache Beam's TestPipeline as a test case +TestPipeline.__test__ = False + + +class SessionizationTest(unittest.TestCase): + """Unit tests for ProcessCustomerSessionDoFn using TestPipeline.""" + + def test_process_customer_session_aggregation(self): + events = [ + (("hh-42", + CustomerInteractionEvent( + event_type=EventType.TRANSACTION.value, + household_key="hh-42", + transaction_id="tx-101", + transaction=TransactionItem( + product_id="prod-A", + quantity=2, + sales_value=20.0, + retail_disc=2.0, + coupon_disc=1.0, + store_id="store-1", + ), + )), 1000), + (("hh-42", + CustomerInteractionEvent( + event_type=EventType.TRANSACTION.value, + household_key="hh-42", + transaction_id="tx-102", + transaction=TransactionItem( + product_id="prod-B", + quantity=1, + sales_value=10.0, + retail_disc=0.0, + coupon_disc=0.0, + store_id="store-1", + ), + )), 1100), + (("hh-42", + CustomerInteractionEvent( + event_type=EventType.COUPON.value, + household_key="hh-42", + transaction_id="tx-101", + coupon=CouponRedemption( + coupon_upc="cp-999", + campaign="fall-sale", + ), + )), 1150), + ] + + with TestPipeline() as p: + results = ( + p + | beam.Create(events) + | beam.Map(lambda x: TimestampedValue(x[0], x[1])) + | beam.WindowInto(Sessions(300)) + | beam.GroupByKey() + | beam.ParDo(ProcessCustomerSessionDoFn()).with_outputs( + TAG_SESSIONS, main="unified_records")) + + def check_unified(records): + assert len(records) == 2 + assert all(isinstance(r, UnifiedTransactionRecord) for r in records) + tx_101 = next(u for u in records if u.transaction_id == "tx-101") + assert tx_101.coupon_upc == "cp-999" + assert tx_101.campaign == "fall-sale" + assert tx_101.sales_value == 20.0 + + tx_102 = next(u for u in records if u.transaction_id == "tx-102") + assert tx_102.coupon_upc is None + + def check_sessions(sessions): + assert len(sessions) == 1 + session = sessions[0] + assert isinstance(session, CustomerSessionProfile) + assert session.household_key == "hh-42" + assert session.total_transactions == 2 + assert session.total_items_purchased == 3 + assert session.total_spend == 30.0 + assert session.total_discount == 3.0 + assert session.coupons_redeemed_count == 1 + assert session.distinct_products_count == 2 + assert "fall-sale" in session.campaigns + assert session.session_duration_sec == 450 + + assert_that(results.unified_records, check_unified, label="CheckUnified") + assert_that(results[TAG_SESSIONS], check_sessions, label="CheckSessions") + + def test_process_customer_session_empty_events(self): + with TestPipeline() as p: + results = ( + p + | beam.Create([("hh-empty", [])]) + | beam.ParDo(ProcessCustomerSessionDoFn()).with_outputs( + TAG_SESSIONS, main="unified_records")) + + assert_that(results.unified_records, equal_to([]), label="CheckEmptyU") + assert_that(results[TAG_SESSIONS], equal_to([]), label="CheckEmptyS") + + def test_process_customer_session_without_coupons(self): + events = [ + (("hh-50", + CustomerInteractionEvent( + event_type=EventType.TRANSACTION.value, + household_key="hh-50", + transaction_id="tx-201", + transaction=TransactionItem( + product_id="prod-X", + quantity=1, + sales_value=5.0, + ), + )), 1000), + ] + + with TestPipeline() as p: + results = ( + p + | beam.Create(events) + | beam.Map(lambda x: TimestampedValue(x[0], x[1])) + | beam.WindowInto(Sessions(300)) + | beam.GroupByKey() + | beam.ParDo(ProcessCustomerSessionDoFn()).with_outputs( + TAG_SESSIONS, main="unified_records")) + + def check_unified(records): + assert len(records) == 1 + assert records[0].coupon_upc is None + assert records[0].campaign is None + + def check_sessions(sessions): + assert len(sessions) == 1 + assert sessions[0].coupons_redeemed_count == 0 + assert sessions[0].campaigns == [] + + assert_that( + results.unified_records, check_unified, label="CheckNoCouponU") + assert_that(results[TAG_SESSIONS], check_sessions, label="CheckNoCouponS") + + def test_process_customer_session_multiple_coupons_same_tx(self): + events = [ + (("hh-60", + CustomerInteractionEvent( + event_type=EventType.TRANSACTION.value, + household_key="hh-60", + transaction_id="tx-301", + transaction=TransactionItem( + product_id="prod-Y", + quantity=2, + sales_value=12.0, + ), + )), 1000), + (("hh-60", + CustomerInteractionEvent( + event_type=EventType.COUPON.value, + household_key="hh-60", + transaction_id="tx-301", + coupon=CouponRedemption( + coupon_upc="cp-1", + campaign="camp-A", + ), + )), 1010), + (("hh-60", + CustomerInteractionEvent( + event_type=EventType.COUPON.value, + household_key="hh-60", + transaction_id="tx-301", + coupon=CouponRedemption( + coupon_upc="cp-2", + campaign="camp-B", + ), + )), 1020), + ] + + with TestPipeline() as p: + results = ( + p + | beam.Create(events) + | beam.Map(lambda x: TimestampedValue(x[0], x[1])) + | beam.WindowInto(Sessions(300)) + | beam.GroupByKey() + | beam.ParDo(ProcessCustomerSessionDoFn()).with_outputs( + TAG_SESSIONS, main="unified_records")) + + def check_unified(records): + assert len(records) == 2 + coupon_upcs = {u.coupon_upc for u in records} + assert coupon_upcs == {"cp-1", "cp-2"} + + def check_sessions(sessions): + assert len(sessions) == 1 + assert sessions[0].coupons_redeemed_count == 2 + assert sessions[0].campaigns == ["camp-A", "camp-B"] + + assert_that(results.unified_records, check_unified, label="CheckMultiU") + assert_that(results[TAG_SESSIONS], check_sessions, label="CheckMultiS") + + +if __name__ == "__main__": + unittest.main() diff --git a/pipelines/cdp/tests/test_sinks.py b/pipelines/cdp/tests/test_sinks.py new file mode 100644 index 00000000..b8801d78 --- /dev/null +++ b/pipelines/cdp/tests/test_sinks.py @@ -0,0 +1,171 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for BigQuery sinks and Storage Write API row formatting.""" + +from datetime import datetime, timezone +import unittest + +import apache_beam as beam +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.utils.timestamp import Timestamp + +from cdp_pipeline.models import ( + CustomerSessionProfile, + DeadLetterRecord, + UnifiedTransactionRecord, +) +from cdp_pipeline.options import MyPipelineOptions +from cdp_pipeline.sinks import ( + _format_deadletter_dict, + _format_session_dict, + _format_unified_dict, + _to_beam_timestamp, + apply_bigquery_sinks, +) + +# Prevent pytest from treating Apache Beam's TestPipeline as a test case +TestPipeline.__test__ = False + + +class SinksTest(unittest.TestCase): + """Unit tests for BigQuery row formatters and sink attachments.""" + + def test_to_beam_timestamp(self): + self.assertIsNone(_to_beam_timestamp(None)) + + beam_ts = Timestamp.of(12345.67) + self.assertEqual(_to_beam_timestamp(beam_ts), beam_ts) + + numeric_ts = 1700000000 + res = _to_beam_timestamp(numeric_ts) + self.assertIsInstance(res, Timestamp) + self.assertEqual(res, Timestamp.of(1700000000.0)) + + dt = datetime(2026, 9, 8, 12, 0, 0, tzinfo=timezone.utc) + res_dt = _to_beam_timestamp(dt) + self.assertIsInstance(res_dt, Timestamp) + self.assertEqual(res_dt, Timestamp.of(dt.timestamp())) + + iso_str = "2026-09-08T12:00:00Z" + res_iso = _to_beam_timestamp(iso_str) + self.assertIsInstance(res_iso, Timestamp) + + invalid_str = "not-a-timestamp" + res_invalid = _to_beam_timestamp(invalid_str) + self.assertIsInstance(res_invalid, Timestamp) + + def test_format_unified_dict_storage_api(self): + record = UnifiedTransactionRecord( + session_id="sess_1", + transaction_id="tx_1", + household_key="hh_1", + product_id="prod_1", + quantity=1, + sales_value=10.0, + store_id="store_1", + retail_disc=0.0, + coupon_discount=0.0, + coupon_match_disc=0.0, + coupon_upc=None, + campaign=None, + day=1, + trans_time="1000", + week_no=1, + event_timestamp="2026-09-08T10:00:00Z", + processed_timestamp="2026-09-08T10:05:00Z", + ) + + formatted_storage = _format_unified_dict(record, use_storage_api=True) + self.assertIsInstance(formatted_storage["event_timestamp"], Timestamp) + self.assertIsInstance(formatted_storage["processed_timestamp"], Timestamp) + self.assertEqual(formatted_storage["transaction_id"], "tx_1") + + formatted_raw = _format_unified_dict(record, use_storage_api=False) + self.assertEqual(formatted_raw["event_timestamp"], "2026-09-08T10:00:00Z") + self.assertEqual(formatted_raw["processed_timestamp"], + "2026-09-08T10:05:00Z") + + def test_format_session_dict_storage_api(self): + record = CustomerSessionProfile( + session_id="sess_1", + household_key="hh_1", + session_start="2026-09-08T10:00:00Z", + session_end="2026-09-08T10:15:00Z", + session_duration_sec=900, + total_transactions=1, + total_items_purchased=1, + total_spend=10.0, + total_discount=0.0, + coupons_redeemed_count=0, + distinct_products_count=1, + campaigns=[], + stores_visited=["store_1"], + processed_timestamp="2026-09-08T10:15:05Z", + ) + + formatted_storage = _format_session_dict(record, use_storage_api=True) + self.assertIsInstance(formatted_storage["session_start"], Timestamp) + self.assertIsInstance(formatted_storage["session_end"], Timestamp) + self.assertIsInstance(formatted_storage["processed_timestamp"], Timestamp) + + formatted_raw = _format_session_dict(record, use_storage_api=False) + self.assertEqual(formatted_raw["session_start"], "2026-09-08T10:00:00Z") + self.assertEqual(formatted_raw["session_end"], "2026-09-08T10:15:00Z") + + def test_format_deadletter_dict_storage_api(self): + record = DeadLetterRecord( + source="transaction", + raw_payload='{"bad": "data"}', + error_message="corrupt", + timestamp="2026-09-08T10:00:00Z", + ) + + formatted_storage = _format_deadletter_dict(record, use_storage_api=True) + self.assertIsInstance(formatted_storage["timestamp"], Timestamp) + self.assertEqual(formatted_storage["source"], "transaction") + + formatted_raw = _format_deadletter_dict(record, use_storage_api=False) + self.assertEqual(formatted_raw["timestamp"], "2026-09-08T10:00:00Z") + + def test_apply_bigquery_sinks_no_project_or_dataset(self): + p = TestPipeline() + pcol = p | "Create" >> beam.Create([]) + options_no_project = MyPipelineOptions([]) + # Should safely return without error + apply_bigquery_sinks(pcol, pcol, pcol, options_no_project) + + def test_apply_bigquery_sinks_graph_construction(self): + options = MyPipelineOptions([ + "--project=test-project", + "--output_dataset=test_dataset", + "--output_table=test_unified", + "--output_sessions_table=test_sessions", + "--deadletter_table=test_dlq", + "--streaming", + ]) + p = TestPipeline(options=options) + dummy_unified = p | "Create Unified" >> beam.Create([]) + dummy_sessions = p | "Create Sessions" >> beam.Create([]) + dummy_deadletters = p | "Create DLQ" >> beam.Create([]) + + apply_bigquery_sinks( + unified_records=dummy_unified, + customer_sessions=dummy_sessions, + all_deadletters=dummy_deadletters, + pipeline_options=options, + ) + + +if __name__ == "__main__": + unittest.main()