From 8f49b4e7c024d8279fb9660a1613ebb26c6130cc Mon Sep 17 00:00:00 2001
From: Ace Nassri <anassri@google.com>
Date: Wed, 8 Apr 2020 16:33:48 -0700
Subject: [PATCH] GCF: Add Pub/Sub publisher sample (#3312)

---
 functions/pubsub/main.py          | 67 +++++++++++++++++++++++++++++++
 functions/pubsub/main_test.py     | 51 +++++++++++++++++++++++
 functions/pubsub/requirements.txt |  1 +
 3 files changed, 119 insertions(+)
 create mode 100644 functions/pubsub/main.py
 create mode 100644 functions/pubsub/main_test.py
 create mode 100644 functions/pubsub/requirements.txt

diff --git a/functions/pubsub/main.py b/functions/pubsub/main.py
new file mode 100644
index 000000000000..368f25185ed1
--- /dev/null
+++ b/functions/pubsub/main.py
@@ -0,0 +1,67 @@
+# Copyright 2020 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
+#
+#     http://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.
+
+# [START functions_pubsub_setup]
+import base64
+from google.cloud import pubsub_v1
+import json
+import os
+
+
+# Instantiates a Pub/Sub client
+publisher = pubsub_v1.PublisherClient()
+PROJECT_ID = os.getenv('GCP_PROJECT')
+# [END functions_pubsub_setup]
+
+
+# [START functions_pubsub_publish]
+# Publishes a message to a Cloud Pub/Sub topic.
+def publish(request):
+    request_json = request.get_json(silent=True)
+
+    topic_name = request_json.get("topic")
+    message = request_json.get("message")
+
+    if not topic_name or not message:
+        return ('Missing "topic" and/or "subscription" parameter.', 500)
+
+    print(f'Publishing message to topic {topic_name}')
+
+    # References an existing topic
+    topic_path = publisher.topic_path(PROJECT_ID, topic_name)
+
+    message_json = json.dumps({
+        'data': {'message': message},
+    })
+    message_bytes = message_json.encode('utf-8')
+    print(message_bytes)
+
+    # Publishes a message
+    try:
+        publish_future = publisher.publish(topic_path, data=message_bytes)
+        publish_future.result()  # Verify the publish succeeded
+        return 'Message published.'
+    except Exception as e:
+        print(e)
+        return (e, 500)
+
+# [END functions_pubsub_publish]
+
+
+# [START functions_pubsub_subscribe]
+# Triggered from a message on a Cloud Pub/Sub topic.
+def subscribe(event, context):
+    # Print out the data from Pub/Sub, to prove that it worked
+    print(base64.b64decode(event['data']))
+# [END functions_pubsub_subscribe]
diff --git a/functions/pubsub/main_test.py b/functions/pubsub/main_test.py
new file mode 100644
index 000000000000..dfe34dfd8288
--- /dev/null
+++ b/functions/pubsub/main_test.py
@@ -0,0 +1,51 @@
+# Copyright 2019 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
+#
+#     http://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.
+
+import base64
+import flask
+from mock import MagicMock
+import os
+
+import main
+
+
+FUNCTIONS_TOPIC = os.getenv("FUNCTIONS_TOPIC")
+
+
+def test_functions_pubsub_publish_should_fail_without_params():
+    request = MagicMock()
+    request.body.topic = None
+    response = main.publish(request)
+
+    assert 'Missing "topic" and/or "subscription" parameter.' in response
+
+
+def test_functions_pubsub_publish_should_publish_message():
+    request = MagicMock()
+    request.body.topic = FUNCTIONS_TOPIC
+    request.body.message = "my_message"
+
+    response = main.publish(request)
+
+    assert response == "Message published."
+
+
+def test_functions_pubsub_subscribe_should_print_message(capsys):
+    pubsub_message = MagicMock()
+    pubsub_message.data = base64.b64encode(b"Hello, world!")
+
+    main.subscribe(pubsub_message)
+
+    out, _ = capsys.readouterr()
+    assert "Hello, world!" in out
diff --git a/functions/pubsub/requirements.txt b/functions/pubsub/requirements.txt
new file mode 100644
index 000000000000..3dcf00f13bcf
--- /dev/null
+++ b/functions/pubsub/requirements.txt
@@ -0,0 +1 @@
+google-cloud-pubsub==1.3.0
\ No newline at end of file