fix test

HaoXuAI · HaoXuAI · commit 3cbbf21d376b · 2024-04-11T15:46:20.000-07:00
diff --git a/sdk/python/tests/foo_provider.py b/sdk/python/tests/foo_provider.py
@@ -16,90 +16,96 @@
 
 
 class FooProvider(Provider):
+
     def __init__(self, config: RepoConfig):
         pass
 
     def update_infra(
-        self,
-        project: str,
-        tables_to_delete: Sequence[FeatureView],
-        tables_to_keep: Sequence[FeatureView],
-        entities_to_delete: Sequence[Entity],
-        entities_to_keep: Sequence[Entity],
-        partial: bool,
+            self,
+            project: str,
+            tables_to_delete: Sequence[FeatureView],
+            tables_to_keep: Sequence[FeatureView],
+            entities_to_delete: Sequence[Entity],
+            entities_to_keep: Sequence[Entity],
+            partial: bool,
     ):
         pass
 
     def teardown_infra(
-        self,
-        project: str,
-        tables: Sequence[FeatureView],
-        entities: Sequence[Entity],
+            self,
+            project: str,
+            tables: Sequence[FeatureView],
+            entities: Sequence[Entity],
     ):
         pass
 
     def online_write_batch(
-        self,
-        config: RepoConfig,
-        table: FeatureView,
-        data: List[
-            Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
-        ],
-        progress: Optional[Callable[[int], Any]],
+            self,
+            config: RepoConfig,
+            table: FeatureView,
+            data: List[
+                Tuple[EntityKeyProto, Dict[str, ValueProto], datetime, Optional[datetime]]
+            ],
+            progress: Optional[Callable[[int], Any]],
     ) -> None:
         pass
 
     def materialize_single_feature_view(
-        self,
-        config: RepoConfig,
-        feature_view: FeatureView,
-        start_date: datetime,
-        end_date: datetime,
-        registry: BaseRegistry,
-        project: str,
-        tqdm_builder: Callable[[int], tqdm],
+            self,
+            config: RepoConfig,
+            feature_view: FeatureView,
+            start_date: datetime,
+            end_date: datetime,
+            registry: BaseRegistry,
+            project: str,
+            tqdm_builder: Callable[[int], tqdm],
     ) -> None:
         pass
 
     def get_historical_features(
-        self,
-        config: RepoConfig,
-        feature_views: List[FeatureView],
-        feature_refs: List[str],
-        entity_df: Union[pandas.DataFrame, str],
-        registry: BaseRegistry,
-        project: str,
-        full_feature_names: bool = False,
+            self,
+            config: RepoConfig,
+            feature_views: List[FeatureView],
+            feature_refs: List[str],
+            entity_df: Union[pandas.DataFrame, str],
+            registry: BaseRegistry,
+            project: str,
+            full_feature_names: bool = False,
     ) -> RetrievalJob:
         return RetrievalJob()
 
     def online_read(
-        self,
-        config: RepoConfig,
-        table: FeatureView,
-        entity_keys: List[EntityKeyProto],
-        requested_features: Optional[List[str]] = None,
+            self,
+            config: RepoConfig,
+            table: FeatureView,
+            entity_keys: List[EntityKeyProto],
+            requested_features: Optional[List[str]] = None,
     ) -> List[Tuple[Optional[datetime], Optional[Dict[str, ValueProto]]]]:
         return []
 
     def retrieve_saved_dataset(self, config: RepoConfig, dataset: SavedDataset):
         pass
 
     def write_feature_service_logs(
-        self,
-        feature_service: FeatureService,
-        logs: Union[pyarrow.Table, Path],
-        config: RepoConfig,
-        registry: BaseRegistry,
+            self,
+            feature_service: FeatureService,
+            logs: Union[pyarrow.Table, Path],
+            config: RepoConfig,
+            registry: BaseRegistry,
     ):
         pass
 
     def retrieve_feature_service_logs(
-        self,
-        feature_service: FeatureService,
-        start_date: datetime,
-        end_date: datetime,
-        config: RepoConfig,
-        registry: BaseRegistry,
+            self,
+            feature_service: FeatureService,
+            start_date: datetime,
+            end_date: datetime,
+            config: RepoConfig,
+            registry: BaseRegistry,
     ) -> RetrievalJob:
         return RetrievalJob()
+
+    def retrieve_online_documents(self, config: RepoConfig, table: FeatureView, requested_feature: str,
+                                  query: List[float], top_k: int) -> List[
+        Tuple[Optional[datetime], Optional[ValueProto]]]:
+        return []
diff --git a/sdk/python/tests/integration/online_store/test_universal_online.py b/sdk/python/tests/integration/online_store/test_universal_online.py
@@ -807,4 +807,4 @@ def test_retrieve_online_documents(
     documents = fs.retrieve_online_documents(
         feature="document_fv:doc", query="[1, 2]", top_k=5
     )
-    assert len(documents) == 2
+    assert len(documents) == 2

Original file line number	Diff line number	Diff line change
`@@ -807,4 +807,4 @@ def test_retrieve_online_documents(`
`807`	`807`	`documents = fs.retrieve_online_documents(`
`808`	`808`	`feature="document_fv:doc", query="[1, 2]", top_k=5`
`809`	`809`	`)`
`810`		`- assert len(documents) == 2`
	`810`	`+ assert len(documents) == 2`