def build_producers(self):
producers = []
project = "testing"
try:
# Create pubsub topics
publisher = PublisherClient()
publisher.create_topic(publisher.topic_path(project, "test10"))
publisher.create_topic(publisher.topic_path(project, "test20"))
# Create pubsub subscriptions
subscriber = SubscriberClient()
subscriber.create_subscription(
subscriber.subscription_path(project, "test10"),
subscriber.topic_path(project, "test10"))
subscriber.create_subscription(
subscriber.subscription_path(project, "test20"),
subscriber.topic_path(project, "test20"))
except AlreadyExists as _:
pass
producers.append(
StreamBuilderFactory.get_producer_pubsub().set_project_id(
project).set_output_topic("test10").build())
producers.append(
StreamBuilderFactory.get_producer_kafka().set_broker_servers(
"localhost:9092").set_output_topic("test1").build())
producers.append(
StreamBuilderFactory.get_producer_kafka().set_broker_servers(
"localhost:9092").set_output_topic("test3").build())
return producers
class PubSubConsumer:
GCLOUD_PUBSUB_PROJECT_ID = 'herbie-app'
def __init__(self):
self._subscriber = SubscriberClient()
def create_subscription(self, topic: str, subscription: str):
topic_path = self._subscriber.topic_path(self.GCLOUD_PUBSUB_PROJECT_ID,
topic)
subscription_path = self._subscriber.subscription_path(
self.GCLOUD_PUBSUB_PROJECT_ID, subscription)
try:
self._subscriber.create_subscription(subscription_path, topic_path)
except AlreadyExists as e:
print(f'Subscription {subscription} already exists.')
pass
def subscribe(self, subscription: str):
project_id = self.GCLOUD_PUBSUB_PROJECT_ID
subscription_path = self._subscriber.subscription_path(
project_id, subscription)
future = self._subscriber.subscribe(subscription_path,
self._subscribe_callback)
print(f'Listening for messages on {subscription_path}')
try:
future.result()
except KeyboardInterrupt:
future.cancel()
def _subscribe_callback(self, message):
print(message.data)
def subscription(topic: str, subscriber: SubscriberClient) -> Iterator[str]:
name = topic.replace("/topics/", "/subscriptions/")
try:
subscriber.create_subscription(name, topic)
delete = True
except AlreadyExists:
delete = False
try:
yield name
finally:
if delete:
subscriber.delete_subscription(name)
def test_submit_pubsub_topic_not_found(
integration_test: IntegrationTest,
publisher: PublisherClient,
subscriber: SubscriberClient,
subscription: str,
topic: str,
):
publisher.delete_topic(topic)
try:
integration_test.assert_accepted_and_queued()
finally:
subscriber.delete_subscription(subscription)
publisher.create_topic(topic)
subscriber.create_subscription(subscription, topic)
integration_test.assert_flushed_and_delivered()
def ensure_topic_and_subscription():
from google.cloud.pubsub_v1 import PublisherClient, SubscriberClient
publisher_client = PublisherClient()
try:
publisher_client.create_topic(
publisher_client.topic_path(PROJECT_ID, TOPIC_NAME))
except AlreadyExists:
pass
subscriber_client = SubscriberClient()
try:
subscriber_client.create_subscription(
subscriber_client.subscription_path(PROJECT_ID, SUBSCRIPTION_NAME),
publisher_client.topic_path(PROJECT_ID, TOPIC_NAME))
except AlreadyExists:
pass
def setUpClass(cls):
super().setUpClass()
cls.project = os.environ.get("PUBSUB_PROJECT", "testing")
cls.credentials = os.environ.get("GOOGLE_APPLICATION_CREDENTIALS",
None)
for subscription in ["test0", "test1"]:
try:
publisher = PublisherClient()
publisher.create_topic(
publisher.topic_path(cls.project, subscription))
except AlreadyExists:
pass
try:
subscriber = SubscriberClient()
subscriber.create_subscription(
subscriber.subscription_path(cls.project, subscription),
subscriber.topic_path(cls.project, subscription))
except AlreadyExists:
pass
def create_subscription(client: pubsub_v1.SubscriberClient, project_name: str,
sub_name: str):
"""Creates a new subscription to a given topic."""
print("using pubsub topic: %s \n" % PUBSUB_TOPIC)
name = client.subscription_path(project_name, sub_name)
topic = client.topic_path(project_name, "new_tweets")
print(name)
print(topic)
try:
subscription: sub_client.types.module.Subscription = client.create_subscription(
name, topic)
except:
subscription = "subscription already exists"
#if subscription already exists, should return ALREADY_EXISTS
print('Subscription {} was created.'.format(subscription))
def subscription_path(subscriber_client: pubsub_v1.SubscriberClient,
topic_path: str) -> Generator[str, None, None]:
subscription_path = subscriber_client.subscription_path(
PROJECT_ID, SUBSCRIPTION_ID)
subscription = subscriber_client.create_subscription(request={
"name": subscription_path,
"topic": topic_path
})
yield subscription.name
try:
subscriber_client.delete_subscription(
request={"subscription": subscription_path})
except NotFound:
pass
def subscription_path(subscriber_client: pubsub_v1.SubscriberClient,
topic_path: str) -> Generator[str, None, None]:
subscription_path = subscriber_client.subscription_path(
PROJECT_ID, SUBSCRIPTION_ID)
try:
subscription = subscriber_client.create_subscription(
request={
"name": subscription_path,
"topic": topic_path
})
yield subscription.name
except AlreadyExists:
yield subscription_path
subscriber_client.delete_subscription(
request={"subscription": subscription_path})
subscriber_client.close()
def proto_subscription(subscriber_client: pubsub_v1.SubscriberClient,
proto_topic: str) -> Generator[str, None, None]:
proto_subscription_path = subscriber_client.subscription_path(
PROJECT_ID, PROTO_SUBSCRIPTION_ID)
try:
proto_subscription = subscriber_client.get_subscription(
request={"subscription": proto_subscription_path})
except NotFound:
proto_subscription = subscriber_client.create_subscription(
request={
"name": proto_subscription_path,
"topic": proto_topic
})
yield proto_subscription.name
subscriber_client.delete_subscription(
request={"subscription": proto_subscription.name})
def test_flush_manager(options: Dict[str, Any], emulator: str, web: str):
print("starting test")
api = CoreV1Api()
# max number of loops to run when waiting for kube actions to complete
max_wait_loops = 20 if options["cluster"] is None else 60
# server has invalid PUBSUB_EMULATOR, so that only flush can deliver messages
static_pvs = options["cluster"] is None
if static_pvs:
create_static_pvs(api)
print("waiting for pods to be healthy")
for _ in range(max_wait_loops):
if all(pod.status.phase == "Running"
for pod in api.list_namespaced_pod("default").items):
break
time.sleep(1)
else:
assert False, "pods did not become healthy"
# create a subscription to the defined topic
print("creating pubsub subscription")
os.environ["PUBSUB_EMULATOR_HOST"] = emulator
sub_client = SubscriberClient()
topic_path = "projects/{project}/topics/{topic}".format(**options)
subscription_path = "projects/{project}/subscriptions/{topic}".format(
**options)
try:
sub_client.create_subscription(subscription_path,
topic_path,
retry=None)
except AlreadyExists:
pass
print("posting message 0")
requests.post(web, headers={
"host": "web"
}, json={
"id": 0
}).raise_for_status()
print("setting up race condition: attached pvc is also deleted")
delete_pvcs(api)
print("setting up race condition: pod unschedulable due to missing pvc")
with pytest.raises(ApiException) as excinfo:
restart_web_pods(api)
assert excinfo.value.reason == "Conflict"
print("posting message 1")
with pytest.raises(requests.exceptions.ConnectionError):
requests.post(web, headers={
"host": "web"
}, json={
"id": 1
}).raise_for_status()
print("starting flush-manager")
# TODO optionally run flush-manager via subprocess.Popen, to ensure testing
# current code and enable code coverage
_sa = kube_resource("kube/flush-manager.sa.yml", **options)
_cluster_role = kube_resource("kube/flush-manager.clusterrole.yml",
**options)
_cluster_role_binding = kube_resource(
"kube/flush-manager.clusterrolebinding.yml", **options)
_role = kube_resource("kube/flush-manager.role.yml", **options)
_role_binding = kube_resource("kube/flush-manager.rolebinding.yml",
**options)
_deploy = kube_resource("kube/flush-manager.deploy.yml", **options)
with _sa, _cluster_role, _cluster_role_binding, _role, _role_binding, _deploy:
print("posting message 2 until successful")
for i in range(max_wait_loops):
try:
requests.post(web, headers={
"host": "web"
}, json={
"id": 2
}).raise_for_status()
except requests.exceptions.ConnectionError:
if i > 0 and static_pvs:
create_static_pvs(api)
else:
break
time.sleep(1)
else:
assert False, "pod did not recover"
# scale to 0 pods
print("scaling web to 0 pods")
AppsV1Api().patch_namespaced_stateful_set_scale(
name="web",
namespace="default",
body=V1StatefulSet(api_version="apps/v1",
kind="StatefulSet",
spec=dict(replicas=0)),
)
# wait for no pvcs
print("waiting for cleanup to complete")
for _ in range(max_wait_loops):
if not api.list_persistent_volume().items:
break
time.sleep(1)
else:
print("pvs were not cleaned up")
assert [] == api.list_persistent_volume().items
# assert jobs and pvcs also deleted
assert [] == list_pvcs(api)
assert [] == BatchV1Api().list_namespaced_job("default").items
# assert received message id 0 and 2
assert [b'{"id": 0}', b'{"id": 2}'] == [
element.message.data
for element in sub_client.pull(subscription_path, 2).received_messages
]
