sdks/python: test enrichment with tecton

Author: mohamedawnallah

sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment.py

@@ -116,3 +116,50 @@ def enrichment_with_vertex_ai_legacy():
         | "Enrich W/ Vertex AI" >> Enrichment(vertex_ai_handler)
         | "Print" >> beam.Map(print))
   # [END enrichment_with_vertex_ai_legacy]
+
+
+def enrichment_with_tecton():
+  # [START enrichment_with_tecton]
+  import apache_beam as beam
+  from apache_beam.transforms.enrichment import Enrichment
+  from apache_beam.transforms.enrichment_handlers.tecton_feature_store import (
+    TectonConnectionConfig, TectonFeaturesRetrievalConfig,
+    TectonFeatureStoreEnrichmentHandler)
+
+  data = [
+      beam.Row(user_id='user_1990251765'),
+      beam.Row(user_id='user_1284832379'),
+      beam.Row(user_id='user_9979340926'),
+  ]
+
+  # Create connection configuration
+  # Your actual Tecton credentials
+  connection_config = TectonConnectionConfig(
+      url='https://explore.tecton.ai',
+      api_key='101142fd7d775e0a1bd9e343cca2a44d'
+  )
+
+  # Create features retrieval configuration
+  # Using your actual fraud detection feature service
+  features_config = TectonFeaturesRetrievalConfig(
+      feature_service_name='fraud_detection_feature_service',
+      entity_id='user_id',
+      workspace_name='prod'
+  )
+
+  # Create the handler with both configurations
+  tecton_handler = TectonFeatureStoreEnrichmentHandler(
+      connection_config=connection_config,
+      features_retrieval_config=features_config
+  )
+
+  with beam.Pipeline() as p:
+    _ = (
+        p
+        | "Create" >> beam.Create(data)
+        | "Enrich W/ Tecton" >> Enrichment(tecton_handler)
+        | "Print" >> beam.Map(print))
+  # [END enrichment_with_tecton]
+
+
+enrichment_with_tecton()

sdks/python/apache_beam/examples/snippets/transforms/elementwise/enrichment_test.py

@@ -28,6 +28,7 @@
   from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_bigtable, \
   enrichment_with_vertex_ai_legacy
   from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_vertex_ai
+  from apache_beam.examples.snippets.transforms.elementwise.enrichment import enrichment_with_tecton
   from apache_beam.io.requestresponse import RequestResponseIO
 except ImportError:
   raise unittest.SkipTest('RequestResponseIO dependencies are not installed')
@@ -60,6 +61,15 @@ def validate_enrichment_with_vertex_ai_legacy():
   return expected
 
 
+def validate_enrichment_with_tecton():
+  expected = '''[START enrichment_with_tecton]
+Row(user_id='user_9979340926', user_transaction_metrics.amount_count_1d_1d=1, user_transaction_metrics.amount_count_3d_1d=3, user_transaction_metrics.amount_count_7d_1d=7, user_transaction_metrics.amount_mean_1d_1d=65.05, user_transaction_metrics.amount_mean_3d_1d=42.72333333333333, user_transaction_metrics.amount_mean_7d_1d=32.955714285714286)
+Row(user_id='user_1990251765', user_transaction_metrics.amount_count_1d_1d=None, user_transaction_metrics.amount_count_3d_1d=2, user_transaction_metrics.amount_count_7d_1d=3, user_transaction_metrics.amount_mean_1d_1d=None, user_transaction_metrics.amount_mean_3d_1d=25.880000000000003, user_transaction_metrics.amount_mean_7d_1d=27.796666666666667)
+Row(user_id='user_1284832379', user_transaction_metrics.amount_count_1d_1d=2, user_transaction_metrics.amount_count_3d_1d=6, user_transaction_metrics.amount_count_7d_1d=12, user_transaction_metrics.amount_mean_1d_1d=111.465, user_transaction_metrics.amount_mean_3d_1d=61.961666666666666, user_transaction_metrics.amount_mean_7d_1d=171.5625)
+  [END enrichment_with_tecton]'''.splitlines()[1:-1]
+  return expected
+
+
 def std_out_to_dict(stdout_lines, row_key):
   output_dict = {}
   for stdout_line in stdout_lines:
@@ -107,6 +117,16 @@ def test_enrichment_with_vertex_ai_legacy(self, mock_stdout):
         std_out_to_dict(output, 'entity_id'),
         std_out_to_dict(expected, 'entity_id'))
 
+  def test_enrichment_with_tecton(self, mock_stdout):
+    enrichment_with_tecton()
+    output = mock_stdout.getvalue().splitlines()
+    expected = validate_enrichment_with_tecton()
+
+    self.assertEqual(len(output), len(expected))
+    self.assertEqual(
+        std_out_to_dict(output, 'user_id'),
+        std_out_to_dict(expected, 'user_id'))
+
 
 if __name__ == '__main__':
   unittest.main()