@@ -709,6 +709,62 @@ def try_get_features():
709709
710710@pytest .mark .timeout (600 )
711711@pytest .mark .run (order = 18 )
712+ def test_basic_ingest_retrieval_multi_entities (client ):
713+ # Set to another project to test ingestion based on current project context
714+ client .set_project (PROJECT_NAME + "_NS1" )
715+ merchant_fs = FeatureSet (
716+ name = "merchant_fs" ,
717+ features = [Feature (name = "merchant_sales" , dtype = ValueType .FLOAT )],
718+ entities = [
719+ Entity ("driver_id" , ValueType .INT64 ),
720+ Entity ("merchant_id" , ValueType .INT64 ),
721+ ],
722+ max_age = Duration (seconds = 3600 ),
723+ )
724+ client .apply (merchant_fs )
725+
726+ N_ROWS = 2
727+ time_offset = datetime .utcnow ().replace (tzinfo = pytz .utc )
728+ merchant_df = pd .DataFrame (
729+ {
730+ "datetime" : [time_offset ] * N_ROWS ,
731+ "driver_id" : [i for i in range (N_ROWS )],
732+ "merchant_id" : [i for i in range (N_ROWS )],
733+ "merchant_sales" : [float (i ) + 0.5 for i in range (N_ROWS )],
734+ }
735+ )
736+ client .ingest ("merchant_fs" , merchant_df , timeout = 600 )
737+
738+ online_request_entity = [
739+ {"driver_id" : 0 , "merchant_id" : 0 },
740+ {"driver_id" : 1 , "merchant_id" : 1 },
741+ ]
742+ online_request_features = ["merchant_sales" ]
743+
744+ def try_get_features ():
745+ response = client .get_online_features (
746+ entity_rows = online_request_entity , feature_refs = online_request_features
747+ )
748+ is_ok = check_online_response ("merchant_sales" , merchant_df , response )
749+ return response , is_ok
750+
751+ online_features_actual = wait_retry_backoff (
752+ retry_fn = try_get_features ,
753+ timeout_secs = 90 ,
754+ timeout_msg = "Timed out trying to get online feature values" ,
755+ )
756+
757+ online_features_expected = {
758+ "driver_id" : [0 , 1 ],
759+ "merchant_id" : [0 , 1 ],
760+ "merchant_sales" : [0.5 , 1.5 ],
761+ }
762+
763+ assert online_features_actual .to_dict () == online_features_expected
764+
765+
766+ @pytest .mark .timeout (600 )
767+ @pytest .mark .run (order = 19 )
712768def test_basic_retrieve_feature_row_missing_fields (client , cust_trans_df ):
713769 feature_refs = ["daily_transactions" , "total_transactions" , "null_values" ]
714770
@@ -756,7 +812,7 @@ def try_get_features():
756812
757813
758814@pytest .mark .timeout (600 )
759- @pytest .mark .run (order = 19 )
815+ @pytest .mark .run (order = 20 )
760816def test_basic_retrieve_feature_row_extra_fields (client , cust_trans_df ):
761817 feature_refs = ["daily_transactions" , "total_transactions" ]
762818 # apply cust_trans_fs and ingest dataframe
@@ -851,7 +907,7 @@ def all_types_dataframe():
851907
852908
853909@pytest .mark .timeout (45 )
854- @pytest .mark .run (order = 20 )
910+ @pytest .mark .run (order = 21 )
855911def test_all_types_register_feature_set_success (client ):
856912 client .set_project (PROJECT_NAME )
857913
@@ -897,7 +953,7 @@ def test_all_types_register_feature_set_success(client):
897953
898954
899955@pytest .mark .timeout (300 )
900- @pytest .mark .run (order = 21 )
956+ @pytest .mark .run (order = 22 )
901957def test_all_types_ingest_success (client , all_types_dataframe ):
902958 # Get all_types feature set
903959 all_types_fs = client .get_feature_set (name = "all_types" )
@@ -907,7 +963,7 @@ def test_all_types_ingest_success(client, all_types_dataframe):
907963
908964
909965@pytest .mark .timeout (90 )
910- @pytest .mark .run (order = 22 )
966+ @pytest .mark .run (order = 23 )
911967def test_all_types_retrieve_online_success (client , all_types_dataframe ):
912968 # Poll serving for feature values until the correct values are returned_float_list
913969 feature_refs = [
0 commit comments