github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/ml/gcp/recommendations_ai_test_it.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Integration tests for Recommendations AI transforms."""
    19  
    20  from __future__ import absolute_import
    21  
    22  import random
    23  import unittest
    24  from datetime import datetime
    25  
    26  import pytest
    27  
    28  import apache_beam as beam
    29  from apache_beam.testing.test_pipeline import TestPipeline
    30  from apache_beam.testing.util import assert_that
    31  from apache_beam.testing.util import equal_to
    32  from apache_beam.testing.util import is_not_empty
    33  
    34  # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
    35  try:
    36    from google.cloud import recommendationengine
    37    from apache_beam.ml.gcp import recommendations_ai
    38  except ImportError:
    39    recommendationengine = None
    40  # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
    41  
    42  GCP_TEST_PROJECT = 'apache-beam-testing'
    43  
    44  CATALOG_ITEM = {
    45      "id": f"aitest-{int(datetime.now().timestamp())}-"
    46      f"{int(random.randint(1,10000))}",
    47      "title": "Sample laptop",
    48      "description": "Indisputably the most fantastic laptop ever created.",
    49      "language_code": "en",
    50      "category_hierarchies": [{
    51          "categories": ["Electronic", "Computers"]
    52      }]
    53  }
    54  
    55  
    56  def extract_id(response):
    57    yield response["id"]
    58  
    59  
    60  def extract_event_type(response):
    61    yield response["event_type"]
    62  
    63  
    64  def extract_prediction(response):
    65    yield response[0]["results"]
    66  
    67  
    68  @pytest.mark.it_postcommit
    69  @unittest.skipIf(
    70      recommendationengine is None,
    71      "Recommendations AI dependencies not installed.")
    72  class RecommendationAIIT(unittest.TestCase):
    73    def test_create_catalog_item(self):
    74  
    75      with TestPipeline(is_integration_test=True) as p:
    76        output = (
    77            p | 'Create data' >> beam.Create([CATALOG_ITEM])
    78            | 'Create CatalogItem' >>
    79            recommendations_ai.CreateCatalogItem(project=GCP_TEST_PROJECT)
    80            | beam.ParDo(extract_id) | beam.combiners.ToList())
    81  
    82        assert_that(output, equal_to([[CATALOG_ITEM["id"]]]))
    83  
    84    def test_create_user_event(self):
    85      USER_EVENT = {"event_type": "page-visit", "user_info": {"visitor_id": "1"}}
    86  
    87      with TestPipeline(is_integration_test=True) as p:
    88        output = (
    89            p | 'Create data' >> beam.Create([USER_EVENT]) | 'Create UserEvent' >>
    90            recommendations_ai.WriteUserEvent(project=GCP_TEST_PROJECT)
    91            | beam.ParDo(extract_event_type) | beam.combiners.ToList())
    92  
    93        assert_that(output, equal_to([[USER_EVENT["event_type"]]]))
    94  
    95    def test_predict(self):
    96      USER_EVENT = {"event_type": "page-visit", "user_info": {"visitor_id": "1"}}
    97  
    98      with TestPipeline(is_integration_test=True) as p:
    99        output = (
   100            p | 'Create data' >> beam.Create([USER_EVENT])
   101            | 'Predict UserEvent' >> recommendations_ai.PredictUserEvent(
   102                project=GCP_TEST_PROJECT, placement_id="recently_viewed_default")
   103            | beam.ParDo(extract_prediction))
   104  
   105        assert_that(output, is_not_empty())
   106  
   107    @classmethod
   108    def tearDownClass(cls):
   109      client = recommendationengine.CatalogServiceClient()
   110      parent = (
   111          f'projects/{GCP_TEST_PROJECT}/locations/'
   112          'global/catalogs/default_catalog')
   113      for item in list(client.list_catalog_items(parent=parent)):
   114        client.delete_catalog_item(
   115            name=f"projects/{GCP_TEST_PROJECT}/locations/global/catalogs/"
   116            f"default_catalog/catalogItems/{item.id}")
   117  
   118  
   119  if __name__ == '__main__':
   120    unittest.main()