Skip to content

Commit 6990f16

Browse files
terryyylimterryyylim
committed
Remove duplicated code in pytest
Signed-off-by: Terence <terencelimxp@gmail.com>
1 parent 404cc36 commit 6990f16

1 file changed

Lines changed: 77 additions & 108 deletions

File tree

‎sdk/python/tests/test_client.py‎

Lines changed: 77 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
import socket
1717
from concurrent import futures
1818
from datetime import datetime, timedelta
19+
from typing import Tuple
1920
from unittest import mock
2021

2122
import grpc
@@ -449,33 +450,8 @@ def test_ingest_dataframe_partition(self, mocked_client, mocker, partitioned_df)
449450
mocker.patch.object(
450451
mocked_client._core_service_stub,
451452
"GetFeatureTable",
452-
return_value=GetFeatureTableResponse(
453-
table=FeatureTableProto(
454-
spec=FeatureTableSpecProto(
455-
name="ingest_featuretable",
456-
max_age=Duration(seconds=3600),
457-
features=[
458-
FeatureSpecProto(
459-
name="dev_feature_float",
460-
value_type=ValueProto.ValueType.FLOAT,
461-
),
462-
FeatureSpecProto(
463-
name="dev_feature_string",
464-
value_type=ValueProto.ValueType.STRING,
465-
),
466-
],
467-
entities=["dev_entity"],
468-
batch_source=DataSourceProto(
469-
type="BATCH_FILE",
470-
file_options=DataSourceProto.FileOptions(
471-
file_format="parquet", file_url="file://feast/*"
472-
),
473-
timestamp_column="datetime",
474-
date_partition_column="datetime_col",
475-
),
476-
),
477-
meta=FeatureTableMetaProto(),
478-
)
453+
return_value=_ingest_test_getfeaturetable_mocked_resp(
454+
"file://feast/*", "datetime_col"
479455
),
480456
)
481457

@@ -486,14 +462,9 @@ def test_ingest_dataframe_partition(self, mocked_client, mocker, partitioned_df)
486462
dest_fpath = os.path.join("feast/")
487463
pq_df = pq.read_table(dest_fpath).to_pandas()
488464

489-
# Format Dataframes before comparing them
490-
partitioned_df.sort_values(by=["dev_feature_float"], inplace=True)
491-
pq_df.sort_values(by=["dev_feature_float"], inplace=True)
492-
pq_df = pq_df.reindex(sorted(pq_df.columns), axis=1)
493-
partitioned_df = partitioned_df.reindex(sorted(partitioned_df.columns), axis=1)
494-
partitioned_df.reset_index(drop=True, inplace=True)
495-
pq_df.reset_index(drop=True, inplace=True)
496-
pq_df["datetime_col"] = pd.to_datetime(pq_df.datetime_col).dt.tz_convert("UTC")
465+
partitioned_df, pq_df = _ingest_test_format_dataframes(
466+
partitioned_df, pq_df, True
467+
)
497468

498469
assert_frame_equal(partitioned_df, pq_df)
499470

@@ -514,33 +485,7 @@ def test_ingest_dataframe_no_partition(
514485
mocker.patch.object(
515486
mocked_client._core_service_stub,
516487
"GetFeatureTable",
517-
return_value=GetFeatureTableResponse(
518-
table=FeatureTableProto(
519-
spec=FeatureTableSpecProto(
520-
name="ingest_featuretable",
521-
max_age=Duration(seconds=3600),
522-
features=[
523-
FeatureSpecProto(
524-
name="dev_feature_float",
525-
value_type=ValueProto.ValueType.FLOAT,
526-
),
527-
FeatureSpecProto(
528-
name="dev_feature_string",
529-
value_type=ValueProto.ValueType.STRING,
530-
),
531-
],
532-
entities=["dev_entity"],
533-
batch_source=DataSourceProto(
534-
type="BATCH_FILE",
535-
file_options=DataSourceProto.FileOptions(
536-
file_format="parquet", file_url="file://feast2/*"
537-
),
538-
timestamp_column="datetime",
539-
),
540-
),
541-
meta=FeatureTableMetaProto(),
542-
)
543-
),
488+
return_value=_ingest_test_getfeaturetable_mocked_resp("file://feast2/*"),
544489
)
545490

546491
mocked_client.set_project("my_project")
@@ -556,15 +501,9 @@ def test_ingest_dataframe_no_partition(
556501
][0]
557502
pq_df = pq.read_table(dest_fpath + single_file).to_pandas()
558503

559-
# Format Dataframes before comparing them
560-
non_partitioned_df.sort_values(by=["dev_feature_float"], inplace=True)
561-
pq_df.sort_values(by=["dev_feature_float"], inplace=True)
562-
pq_df = pq_df.reindex(sorted(pq_df.columns), axis=1)
563-
non_partitioned_df = non_partitioned_df.reindex(
564-
sorted(non_partitioned_df.columns), axis=1
504+
non_partitioned_df, pq_df = _ingest_test_format_dataframes(
505+
non_partitioned_df, pq_df
565506
)
566-
non_partitioned_df.reset_index(drop=True, inplace=True)
567-
pq_df.reset_index(drop=True, inplace=True)
568507

569508
assert_frame_equal(non_partitioned_df, pq_df)
570509

@@ -583,33 +522,8 @@ def test_ingest_csv(self, mocked_client, mocker):
583522
mocker.patch.object(
584523
mocked_client._core_service_stub,
585524
"GetFeatureTable",
586-
return_value=GetFeatureTableResponse(
587-
table=FeatureTableProto(
588-
spec=FeatureTableSpecProto(
589-
name="ingest_featuretable",
590-
max_age=Duration(seconds=3600),
591-
features=[
592-
FeatureSpecProto(
593-
name="dev_feature_float",
594-
value_type=ValueProto.ValueType.FLOAT,
595-
),
596-
FeatureSpecProto(
597-
name="dev_feature_string",
598-
value_type=ValueProto.ValueType.STRING,
599-
),
600-
],
601-
entities=["dev_entity"],
602-
batch_source=DataSourceProto(
603-
type="BATCH_FILE",
604-
file_options=DataSourceProto.FileOptions(
605-
file_format="parquet", file_url="file://feast3/*"
606-
),
607-
timestamp_column="datetime",
608-
date_partition_column="datetime_col",
609-
),
610-
),
611-
meta=FeatureTableMetaProto(),
612-
)
525+
return_value=_ingest_test_getfeaturetable_mocked_resp(
526+
"file://feast3/*", "datetime_col"
613527
),
614528
)
615529

@@ -627,17 +541,9 @@ def test_ingest_csv(self, mocked_client, mocker):
627541
dest_fpath = os.path.join("feast3/")
628542
pq_df = pq.read_table(dest_fpath).to_pandas()
629543

630-
# Format Dataframes before comparing them
631-
partitioned_df.sort_values(by=["dev_feature_float"], inplace=True)
632-
pq_df.sort_values(by=["dev_feature_float"], inplace=True)
633-
pq_df = pq_df.reindex(sorted(pq_df.columns), axis=1)
634-
partitioned_df = partitioned_df.reindex(sorted(partitioned_df.columns), axis=1)
635-
partitioned_df.reset_index(drop=True, inplace=True)
636-
pq_df.reset_index(drop=True, inplace=True)
637-
partitioned_df["datetime_col"] = pd.to_datetime(
638-
partitioned_df.datetime_col
639-
).dt.tz_convert("UTC")
640-
pq_df["datetime_col"] = pd.to_datetime(pq_df.datetime_col).dt.tz_convert("UTC")
544+
partitioned_df, pq_df = _ingest_test_format_dataframes(
545+
partitioned_df, pq_df, True
546+
)
641547

642548
assert_frame_equal(partitioned_df, pq_df)
643549

@@ -710,3 +616,66 @@ def test_no_auth_sent_when_auth_disabled(
710616
):
711617
client = Client(core_url=f"localhost:{insecure_core_server_that_blocks_auth}")
712618
client.list_feature_tables()
619+
620+
621+
def _ingest_test_getfeaturetable_mocked_resp(
622+
file_url: str, date_partition_col: str = None
623+
):
624+
return GetFeatureTableResponse(
625+
table=FeatureTableProto(
626+
spec=FeatureTableSpecProto(
627+
name="ingest_featuretable",
628+
max_age=Duration(seconds=3600),
629+
features=[
630+
FeatureSpecProto(
631+
name="dev_feature_float", value_type=ValueProto.ValueType.FLOAT,
632+
),
633+
FeatureSpecProto(
634+
name="dev_feature_string",
635+
value_type=ValueProto.ValueType.STRING,
636+
),
637+
],
638+
entities=["dev_entity"],
639+
batch_source=DataSourceProto(
640+
file_options=DataSourceProto.FileOptions(
641+
file_format="parquet", file_url=file_url
642+
),
643+
timestamp_column="datetime",
644+
date_partition_column=date_partition_col
645+
if date_partition_col is not None
646+
else None,
647+
),
648+
),
649+
meta=FeatureTableMetaProto(),
650+
)
651+
)
652+
653+
654+
def _ingest_test_format_dataframes(
655+
partitioned_df: pd.DataFrame, pq_df: pd.DataFrame, with_partitions: bool = False
656+
) -> Tuple[pd.DataFrame, pd.DataFrame]:
657+
"""
658+
Format Dataframes before comparing them through assertion.
659+
660+
Args:
661+
partitioned_df: DataFrame from pytest fixture
662+
pq_df: DataFrame from parquet files
663+
with_partitions: Flag to indicate if data has been partitioned
664+
665+
Returns:
666+
Formatted DataFrames for comparison
667+
"""
668+
partitioned_df.sort_values(by=["dev_feature_float"], inplace=True)
669+
pq_df.sort_values(by=["dev_feature_float"], inplace=True)
670+
pq_df = pq_df.reindex(sorted(pq_df.columns), axis=1)
671+
partitioned_df = partitioned_df.reindex(sorted(partitioned_df.columns), axis=1)
672+
partitioned_df.reset_index(drop=True, inplace=True)
673+
pq_df.reset_index(drop=True, inplace=True)
674+
675+
if with_partitions:
676+
partitioned_df["datetime_col"] = pd.to_datetime(
677+
partitioned_df.datetime_col
678+
).dt.tz_convert("UTC")
679+
pq_df["datetime_col"] = pd.to_datetime(pq_df.datetime_col).dt.tz_convert("UTC")
680+
681+
return partitioned_df, pq_df

0 commit comments

Comments
 (0)