Skip to content

Commit ce9d461

Browse files
khorshuhengkhorshuheng
committed
Add integration test for dataproc launcher
Signed-off-by: Khor Shu Heng <khor.heng@gojek.com>
1 parent 05bef7e commit ce9d461

6 files changed

Lines changed: 178 additions & 2 deletions

File tree

‎infra/scripts/test-integration.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,4 +4,4 @@ python -m pip install --upgrade pip setuptools wheel
44
make install-python
55
python -m pip install -qr tests/requirements.txt
66

7-
pytest tests/integration/
7+
pytest tests/integration --dataproc-cluster-name feast-e2e --dataproc-project kf-feast --dataproc-region us-central1 --staging-location gs://feast-templocation-kf-feast

‎tests/integration/conftest.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,7 @@
11
def pytest_addoption(parser):
2-
pass
2+
parser.addoption("--dataproc-cluster-name", action="store")
3+
parser.addoption("--dataproc-region", action="store")
4+
parser.addoption("--dataproc-project", action="store")
5+
parser.addoption("--dataproc-staging-location", action="store")
6+
parser.addoption("--redis-url", action="store")
7+
parser.addoption("--redis-cluster", action="store_true")

‎tests/integration/fixtures/__init__.py‎

Whitespace-only changes.

‎tests/integration/fixtures/job_parameters.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
import tempfile
2+
import uuid
3+
from datetime import datetime
4+
from os import path
5+
from urllib.parse import urlparse
6+
7+
import numpy as np
8+
import pandas as pd
9+
import pytest
10+
from google.cloud import storage
11+
from pytz import utc
12+
13+
from feast.pyspark.abc import RetrievalJobParameters
14+
15+
16+
@pytest.fixture(scope="module")
17+
def customer_entity() -> pd.DataFrame:
18+
return pd.DataFrame(
19+
np.array([[1001, datetime(year=2020, month=9, day=1, tzinfo=utc)]]),
20+
columns=["customer_id", "event_timestamp"],
21+
)
22+
23+
24+
@pytest.fixture(scope="module")
25+
def customer_feature() -> pd.DataFrame:
26+
return pd.DataFrame(
27+
np.array(
28+
[
29+
[
30+
1001,
31+
100.0,
32+
datetime(year=2020, month=9, day=1, tzinfo=utc),
33+
datetime(year=2020, month=9, day=1, tzinfo=utc),
34+
],
35+
]
36+
),
37+
columns=[
38+
"customer_id",
39+
"total_transactions",
40+
"event_timestamp",
41+
"created_timestamp",
42+
],
43+
)
44+
45+
46+
def upload_dataframe_to_gcs_as_parquet(df: pd.DataFrame, staging_location: str):
47+
gcs_client = storage.Client()
48+
staging_location_uri = urlparse(staging_location)
49+
staging_bucket = staging_location_uri.netloc
50+
remote_path = staging_location_uri.path.lstrip("/")
51+
gcs_bucket = gcs_client.get_bucket(staging_bucket)
52+
temp_dir = str(uuid.uuid4())
53+
df_remote_path = path.join(remote_path, temp_dir)
54+
blob = gcs_bucket.blob(df_remote_path)
55+
with tempfile.NamedTemporaryFile() as df_local_path:
56+
df.to_parquet(df_local_path.name)
57+
blob.upload_from_filename(df_local_path.name)
58+
return path.join(staging_location, df_remote_path)
59+
60+
61+
def new_retrieval_job_params(
62+
entity_source_uri: str, feature_source_uri: str, destination_uri: str
63+
) -> RetrievalJobParameters:
64+
entity_source = {
65+
"file": {
66+
"format": "parquet",
67+
"path": entity_source_uri,
68+
"event_timestamp_column": "event_timestamp",
69+
}
70+
}
71+
72+
feature_tables_sources = [
73+
{
74+
"file": {
75+
"format": "parquet",
76+
"path": feature_source_uri,
77+
"event_timestamp_column": "event_timestamp",
78+
"created_timestamp_column": "created_timestamp",
79+
}
80+
}
81+
]
82+
83+
feature_tables = [
84+
{
85+
"name": "customer_transactions",
86+
"entities": [{"name": "customer", "type": "int32"}],
87+
}
88+
]
89+
90+
destination = {"format": "parquet", "path": destination_uri}
91+
92+
return RetrievalJobParameters(
93+
feature_tables=feature_tables,
94+
feature_tables_sources=feature_tables_sources,
95+
entity_source=entity_source,
96+
destination=destination,
97+
)
98+
99+
100+
@pytest.fixture(scope="module")
101+
def dataproc_retrieval_job_params(
102+
pytestconfig, customer_entity, customer_feature
103+
) -> RetrievalJobParameters:
104+
staging_location = pytestconfig.getoption("--dataproc-staging-location")
105+
entity_source_uri = upload_dataframe_to_gcs_as_parquet(
106+
customer_entity, staging_location
107+
)
108+
feature_source_uri = upload_dataframe_to_gcs_as_parquet(
109+
customer_feature, staging_location
110+
)
111+
destination_uri = path.join(staging_location, str(uuid.uuid4()))
112+
113+
return new_retrieval_job_params(
114+
entity_source_uri, feature_source_uri, destination_uri
115+
)

‎tests/integration/fixtures/launchers.py‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
import pytest
2+
3+
from feast.pyspark.launchers.gcloud import DataprocClusterLauncher
4+
5+
6+
@pytest.fixture
7+
def dataproc_launcher(pytestconfig) -> DataprocClusterLauncher:
8+
cluster_name = pytestconfig.getoption("--dataproc-cluster-name")
9+
region = pytestconfig.getoption("--dataproc-region")
10+
project_id = pytestconfig.getoption("--dataproc-project")
11+
staging_location = pytestconfig.getoption("--dataproc-staging-location")
12+
return DataprocClusterLauncher(
13+
cluster_name=cluster_name,
14+
staging_location=staging_location,
15+
region=region,
16+
project_id=project_id,
17+
)

‎tests/integration/test_launchers.py‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
from time import sleep
2+
3+
from feast.pyspark.abc import RetrievalJobParameters, SparkJobStatus
4+
from feast.pyspark.launchers.gcloud import DataprocClusterLauncher
5+
6+
from .fixtures.job_parameters import customer_entity # noqa: F401
7+
from .fixtures.job_parameters import customer_feature # noqa: F401
8+
from .fixtures.job_parameters import dataproc_retrieval_job_params # noqa: F401
9+
from .fixtures.launchers import dataproc_launcher # noqa: F401
10+
11+
12+
def test_dataproc_job_api(
13+
dataproc_launcher: DataprocClusterLauncher, # noqa: F811
14+
dataproc_retrieval_job_params: RetrievalJobParameters, # noqa: F811
15+
):
16+
job = dataproc_launcher.historical_feature_retrieval(dataproc_retrieval_job_params)
17+
job_id = job.get_id()
18+
retrieved_job = dataproc_launcher.get_job_by_id(job_id)
19+
assert retrieved_job.get_id() == job_id
20+
status = retrieved_job.get_status()
21+
assert status in [
22+
SparkJobStatus.STARTING,
23+
SparkJobStatus.IN_PROGRESS,
24+
SparkJobStatus.COMPLETED,
25+
]
26+
active_job_ids = [
27+
job.get_id() for job in dataproc_launcher.list_jobs(include_terminated=False)
28+
]
29+
assert job_id in active_job_ids
30+
retrieved_job.cancel()
31+
sleep(10)
32+
active_job_ids = [
33+
job.get_id() for job in dataproc_launcher.list_jobs(include_terminated=False)
34+
]
35+
assert job_id not in active_job_ids
36+
all_job_ids = [
37+
job.get_id() for job in dataproc_launcher.list_jobs(include_terminated=True)
38+
]
39+
assert job_id in all_job_ids

0 commit comments

Comments
 (0)