github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/ml/gcp/recommendations_ai_test.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  """Unit tests for Recommendations AI transforms."""
    19  
    20  from __future__ import absolute_import
    21  
    22  import unittest
    23  
    24  import mock
    25  
    26  import apache_beam as beam
    27  from apache_beam.metrics import MetricsFilter
    28  
    29  # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
    30  try:
    31    from google.cloud import recommendationengine
    32    from apache_beam.ml.gcp import recommendations_ai
    33  except ImportError:
    34    recommendationengine = None
    35  # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports
    36  
    37  
    38  @unittest.skipIf(
    39      recommendationengine is None,
    40      "Recommendations AI dependencies not installed.")
    41  class RecommendationsAICatalogItemTest(unittest.TestCase):
    42    def setUp(self):
    43      self._mock_client = mock.Mock()
    44      self._mock_client.create_catalog_item.return_value = (
    45          recommendationengine.CatalogItem())
    46      self.m2 = mock.Mock()
    47      self.m2.result.return_value = None
    48      self._mock_client.import_catalog_items.return_value = self.m2
    49  
    50      self._catalog_item = {
    51          "id": "12345",
    52          "title": "Sample laptop",
    53          "description": "Indisputably the most fantastic laptop ever created.",
    54          "language_code": "en",
    55          "category_hierarchies": [{
    56              "categories": ["Electronic", "Computers"]
    57          }]
    58      }
    59  
    60    def test_CreateCatalogItem(self):
    61      expected_counter = 1
    62      with mock.patch.object(recommendations_ai,
    63                             'get_recommendation_catalog_client',
    64                             return_value=self._mock_client):
    65        p = beam.Pipeline()
    66  
    67        _ = (
    68            p | "Create data" >> beam.Create([self._catalog_item])
    69            | "Create CatalogItem" >>
    70            recommendations_ai.CreateCatalogItem(project="test"))
    71  
    72        result = p.run()
    73        result.wait_until_finish()
    74  
    75        read_filter = MetricsFilter().with_name('api_calls')
    76        query_result = result.metrics().query(read_filter)
    77        if query_result['counters']:
    78          read_counter = query_result['counters'][0]
    79          self.assertTrue(read_counter.result == expected_counter)
    80  
    81    def test_ImportCatalogItems(self):
    82      expected_counter = 1
    83      with mock.patch.object(recommendations_ai,
    84                             'get_recommendation_catalog_client',
    85                             return_value=self._mock_client):
    86        p = beam.Pipeline()
    87  
    88        _ = (
    89            p | "Create data" >> beam.Create([
    90                (self._catalog_item["id"], self._catalog_item),
    91                (self._catalog_item["id"], self._catalog_item)
    92            ]) | "Create CatalogItems" >>
    93            recommendations_ai.ImportCatalogItems(project="test"))
    94  
    95        result = p.run()
    96        result.wait_until_finish()
    97  
    98        read_filter = MetricsFilter().with_name('api_calls')
    99        query_result = result.metrics().query(read_filter)
   100        if query_result['counters']:
   101          read_counter = query_result['counters'][0]
   102          self.assertTrue(read_counter.result == expected_counter)
   103  
   104  
   105  @unittest.skipIf(
   106      recommendationengine is None,
   107      "Recommendations AI dependencies not installed.")
   108  class RecommendationsAIUserEventTest(unittest.TestCase):
   109    def setUp(self):
   110      self._mock_client = mock.Mock()
   111      self._mock_client.write_user_event.return_value = (
   112          recommendationengine.UserEvent())
   113      self.m2 = mock.Mock()
   114      self.m2.result.return_value = None
   115      self._mock_client.import_user_events.return_value = self.m2
   116  
   117      self._user_event = {
   118          "event_type": "page-visit", "user_info": {
   119              "visitor_id": "1"
   120          }
   121      }
   122  
   123    def test_CreateUserEvent(self):
   124      expected_counter = 1
   125      with mock.patch.object(recommendations_ai,
   126                             'get_recommendation_user_event_client',
   127                             return_value=self._mock_client):
   128        p = beam.Pipeline()
   129  
   130        _ = (
   131            p | "Create data" >> beam.Create([self._user_event])
   132            | "Create UserEvent" >>
   133            recommendations_ai.WriteUserEvent(project="test"))
   134  
   135        result = p.run()
   136        result.wait_until_finish()
   137  
   138        read_filter = MetricsFilter().with_name('api_calls')
   139        query_result = result.metrics().query(read_filter)
   140        if query_result['counters']:
   141          read_counter = query_result['counters'][0]
   142          self.assertTrue(read_counter.result == expected_counter)
   143  
   144    def test_ImportUserEvents(self):
   145      expected_counter = 1
   146      with mock.patch.object(recommendations_ai,
   147                             'get_recommendation_user_event_client',
   148                             return_value=self._mock_client):
   149        p = beam.Pipeline()
   150  
   151        _ = (
   152            p | "Create data" >> beam.Create([
   153                (self._user_event["user_info"]["visitor_id"], self._user_event),
   154                (self._user_event["user_info"]["visitor_id"], self._user_event)
   155            ]) | "Create UserEvents" >>
   156            recommendations_ai.ImportUserEvents(project="test"))
   157  
   158        result = p.run()
   159        result.wait_until_finish()
   160  
   161        read_filter = MetricsFilter().with_name('api_calls')
   162        query_result = result.metrics().query(read_filter)
   163        if query_result['counters']:
   164          read_counter = query_result['counters'][0]
   165          self.assertTrue(read_counter.result == expected_counter)
   166  
   167  
   168  @unittest.skipIf(
   169      recommendationengine is None,
   170      "Recommendations AI dependencies not installed.")
   171  class RecommendationsAIPredictTest(unittest.TestCase):
   172    def setUp(self):
   173      self._mock_client = mock.Mock()
   174      self._mock_client.predict.return_value = [
   175          recommendationengine.PredictResponse()
   176      ]
   177  
   178      self._user_event = {
   179          "event_type": "page-visit", "user_info": {
   180              "visitor_id": "1"
   181          }
   182      }
   183  
   184    def test_Predict(self):
   185      expected_counter = 1
   186      with mock.patch.object(recommendations_ai,
   187                             'get_recommendation_prediction_client',
   188                             return_value=self._mock_client):
   189        p = beam.Pipeline()
   190  
   191        _ = (
   192            p | "Create data" >> beam.Create([self._user_event])
   193            | "Prediction UserEvents" >> recommendations_ai.PredictUserEvent(
   194                project="test", placement_id="recently_viewed_default"))
   195  
   196        result = p.run()
   197        result.wait_until_finish()
   198  
   199        read_filter = MetricsFilter().with_name('api_calls')
   200        query_result = result.metrics().query(read_filter)
   201        if query_result['counters']:
   202          read_counter = query_result['counters'][0]
   203          self.assertTrue(read_counter.result == expected_counter)
   204  
   205  
   206  if __name__ == '__main__':
   207    unittest.main()