def test_model_predict(self): predictions = [{'output_1': [0.901], 'output_2': [0.997]}] builder = http.RequestMockBuilder({ 'ml.projects.predict': (None, self._make_response_body(predictions, successful=True)) }) resource = discovery.build('ml', 'v1', http=http.HttpMock( self._discovery_testdata_dir, {'status': http_client.OK}), requestBuilder=builder) with mock.patch('googleapiclient.discovery.' 'build') as response_mock: response_mock.side_effect = lambda service, version: resource inference_spec_type = model_spec_pb2.InferenceSpecType( ai_platform_prediction_model_spec=model_spec_pb2. AIPlatformPredictionModelSpec( project_id='test-project', model_name='test-model', )) prediction_log = prediction_log_pb2.PredictionLog() prediction_log.predict_log.response.outputs['output_1'].CopyFrom( tf.make_tensor_proto(values=[0.901], dtype=tf.double, shape=(1, 1))) prediction_log.predict_log.response.outputs['output_2'].CopyFrom( tf.make_tensor_proto(values=[0.997], dtype=tf.double, shape=(1, 1))) self._set_up_pipeline(inference_spec_type) assert_that(self.pcoll, equal_to([prediction_log])) self._run_inference_with_beam()
def test_exception_raised_when_response_body_contains_error_entry(self): error_msg = 'Base64 decode failed.' builder = http.RequestMockBuilder({ 'ml.projects.predict': (None, self._make_response_body(error_msg, successful=False)) }) resource = discovery.build('ml', 'v1', http=http.HttpMock( self._discovery_testdata_dir, {'status': http_client.OK}), requestBuilder=builder) with mock.patch('googleapiclient.discovery.' 'build') as response_mock: response_mock.side_effect = lambda service, version: resource inference_spec_type = model_spec_pb2.InferenceSpecType( ai_platform_prediction_model_spec=model_spec_pb2. AIPlatformPredictionModelSpec( project_id='test-project', model_name='test-model', )) try: self._set_up_pipeline(inference_spec_type) self._run_inference_with_beam() except ValueError as exc: actual_error_msg = str(exc) self.assertTrue(actual_error_msg.startswith(error_msg)) else: self.fail('Test was expected to throw ValueError exception')
def test_build_http(self): """Verify set user agent sets the user agent correctly.""" http_mock = http.HttpMock() h = base.http_helpers.build_http(http=http_mock) _ = h.request('http://test.foo', 'GET') self.assertTrue( forseti_security.__package_name__ in h.headers.get('user-agent'))
def mock_http_response(response, status='200'): """Set the mock response to an http request.""" http_mock = http.HttpMock() http_mock.response_headers = { 'status': status, 'content-type': 'application/json', } http_mock.data = response _base_repository.LOCAL_THREAD.http = http_mock
def test_build_http_multiple(self): """Verify set user agent sets the user agent only once.""" http_mock = http.HttpMock() h = base.http_helpers.build_http(http=http_mock) for _ in range(5): h = base.http_helpers.build_http(http=h) _ = h.request('http://test.foo', 'GET') user_agent = h.headers.get('user-agent') forseti_agent_count = user_agent.count( forseti_security.__package_name__) self.assertEqual(1, forseti_agent_count)
def test_set_user_agent(self): """Verify set user agent sets the user agent correctly.""" credentials = self.get_test_credential() self.assertEqual('', credentials.user_agent) base._set_user_agent(credentials) self.assertTrue( forseti_security.__package_name__ in credentials.user_agent) http_mock = http.HttpMock() credentials.authorize(http_mock) # The user-agent header is set during the request self.assertEqual(None, http_mock.headers) _ = http_mock.request('http://test.foo', 'GET') self.assertTrue(forseti_security.__package_name__ in http_mock.headers.get('user-agent'))
def test_can_format_requests(self): predictions = [{ 'output_1': [0.901], 'output_2': [0.997] }] * len(self._predict_examples) builder = http.RequestMockBuilder({ 'ml.projects.predict': (None, self._make_response_body(predictions, successful=True)) }) resource = discovery.build( 'ml', 'v1', http=http.HttpMock(self._discovery_testdata_dir, {'status': http_client.OK}), requestBuilder=builder) with mock.patch('googleapiclient.discovery.' 'build') as response_mock: response_mock.side_effect = lambda service, version: resource inference_spec_type = model_spec_pb2.InferenceSpecType( ai_platform_prediction_model_spec=model_spec_pb2 .AIPlatformPredictionModelSpec( project_id='test-project', model_name='test-model', )) example = text_format.Parse( """ features { feature { key: "x_bytes" value { bytes_list { value: ["ASa8asdf"] }}} feature { key: "x" value { bytes_list { value: "JLK7ljk3" }}} feature { key: "y" value { int64_list { value: [1, 2] }}} feature { key: "z" value { float_list { value: [4.5, 5, 5.5] }}} } """, tf.train.Example()) self.pipeline = self._make_beam_pipeline() self.pcoll = ( self.pipeline | 'CreateExamples' >> beam.Create([example]) | 'RunInference' >> run_inference.RunInferenceImpl(inference_spec_type)) self._run_inference_with_beam()
def test_model_predict(self, keyed_input: bool): predictions = [{ 'output_1': [0.901], 'output_2': [0.997] }] * len(self._predict_examples) builder = http.RequestMockBuilder({ 'ml.projects.predict': (None, self._make_response_body(predictions, successful=True)) }) resource = discovery.build( 'ml', 'v1', http=http.HttpMock(self._discovery_testdata_dir, {'status': http_client.OK}), requestBuilder=builder) with mock.patch('googleapiclient.discovery.' 'build') as response_mock: response_mock.side_effect = lambda service, version: resource inference_spec_type = model_spec_pb2.InferenceSpecType( ai_platform_prediction_model_spec=model_spec_pb2 .AIPlatformPredictionModelSpec( project_id='test-project', model_name='test-model', )) expected = [] for example in self._predict_examples: prediction_log = prediction_log_pb2.PredictionLog() predict_log = prediction_log.predict_log input_tensor_proto = predict_log.request.inputs['inputs'] input_tensor_proto.dtype = tf.string.as_datatype_enum input_tensor_proto.tensor_shape.dim.add().size = 1 input_tensor_proto.string_val.append(example.SerializeToString()) predict_log.response.outputs['output_1'].CopyFrom( tf.make_tensor_proto(values=[0.901], dtype=tf.double, shape=(1, 1))) predict_log.response.outputs['output_2'].CopyFrom( tf.make_tensor_proto(values=[0.997], dtype=tf.double, shape=(1, 1))) expected.append(prediction_log) self._set_up_pipeline(inference_spec_type, keyed_input) assert_that(self.pcoll, equal_to(expected)) self._run_inference_with_beam()