Integration testing of a Beam Pipeline - apache-beam

Question

We have a Beam pipeline that is written using the Python SDK. The pipeline is reading from an ordered Pub/Sub topic, doing some small transformations and then inserting into ElasticSearch.
We want to write some integration tests that will cover the entire flow, reading from Pub/Sub and using the ES client, to make some assertions over what was inserted as a result. Since the job is streaming, it's a bit tricky to test it inside tests. Here is an example of how I do it now:
#pytest.fixture
def test_options():
setup_options = SetupOptions()
standard_options = setup_options.view_as(StandardOptions)
standard_options.streaming = True
standard_options.runner = "DirectRunner"
return standard_options
def test_reading_from_pubsub_to_elastic(test_options, customer_es_client, pubsub_publisher, user_events):
expected_ids = ["1_3", "1_9", "1_4", "1_6", "1_10", "1_8", "1_5", "1_7", "1_2", "1_1"]
p = HelpderClassToInitPipeline(pipeline=TestPipeline(blocking=False, options=test_options))
pipe = p.create_pipeline(project=PROJECT, local_config=LocalConfig(run_local=False))
pipe_result = pipe.run()
futures = []
for e in user_events:
futures.append(
pubsub_publisher.publish(
topic=f"topic_",
data=json.dumps(e).encode(),
)
)
for f in futures:
f.result()
time.sleep(10) # wait for the pipeline to consume events
result = customer_es_client.search(index="customer", body={"_source": ["_id"], "query": {"match_all": {}}})
ids = [d["_id"] for d in result["hits"]["hits"]]
assert ids.sort() == expected_ids.sort()
pipe_result.cancel()
We are using pytest as a test runner, and we run the inside the CI. Basically we provision Pub/Sub topics in CI, then run the tests, and then we clean up everything.
You can see that time.sleep call that is needed for the pipeline to consume and store so that we could query ES. Not sure if this is the best practice.
Any advices how can I improve this?