Skip to content

Commit

Permalink
Change test_cli test to keep only specs
Browse files Browse the repository at this point in the history
Signed-off-by: Felix Wang <wangfelix98@gmail.com>
  • Loading branch information
felixwang9817 committed Jan 7, 2022
1 parent 349b769 commit 4063d3d
Show file tree
Hide file tree
Showing 2 changed files with 19 additions and 32 deletions.
35 changes: 6 additions & 29 deletions sdk/python/feast/registry.py
Original file line number Diff line number Diff line change
Expand Up @@ -708,69 +708,46 @@ def teardown(self):
"""Tears down (removes) the registry."""
self._registry_store.teardown()

def to_dict(self, project: str, spec_only: bool = False) -> Dict[str, List[Any]]:
def to_dict(self, project: str) -> Dict[str, List[Any]]:
"""Returns a dictionary representation of the registry contents for the specified project.
For each list in the dictionary, the elements are sorted by name, so this
method can be used to compare two registries.
Args:
project: Feast project to convert to a dict
spec_only: If True, return only specs, and exclude metadata.
"""
registry_dict = defaultdict(list)

for entity in sorted(
self.list_entities(project=project), key=lambda entity: entity.name
):
registry_dict["entities"].append(
MessageToDict(
entity.to_proto().spec if spec_only else entity.to_proto()
)
)
registry_dict["entities"].append(MessageToDict(entity.to_proto()))
for feature_view in sorted(
self.list_feature_views(project=project),
key=lambda feature_view: feature_view.name,
):
registry_dict["featureViews"].append(
MessageToDict(
feature_view.to_proto().spec
if spec_only
else feature_view.to_proto()
)
)
registry_dict["featureViews"].append(MessageToDict(feature_view.to_proto()))
for feature_service in sorted(
self.list_feature_services(project=project),
key=lambda feature_service: feature_service.name,
):
registry_dict["featureServices"].append(
MessageToDict(
feature_service.to_proto().spec
if spec_only
else feature_service.to_proto()
)
MessageToDict(feature_service.to_proto())
)
for on_demand_feature_view in sorted(
self.list_on_demand_feature_views(project=project),
key=lambda on_demand_feature_view: on_demand_feature_view.name,
):
registry_dict["onDemandFeatureViews"].append(
MessageToDict(
on_demand_feature_view.to_proto().spec
if spec_only
else on_demand_feature_view.to_proto()
)
MessageToDict(on_demand_feature_view.to_proto())
)
for request_feature_view in sorted(
self.list_request_feature_views(project=project),
key=lambda request_feature_view: request_feature_view.name,
):
registry_dict["requestFeatureViews"].append(
MessageToDict(
request_feature_view.to_proto().spec
if spec_only
else request_feature_view.to_proto()
)
MessageToDict(request_feature_view.to_proto())
)
return registry_dict

Expand Down
16 changes: 13 additions & 3 deletions sdk/python/tests/integration/registration/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,13 @@ def test_universal_cli(test_repo_config) -> None:

# Store registry contents, to be compared later.
fs = FeatureStore(repo_path=str(repo_path))
registry_dict = fs.registry.to_dict(project=project, spec_only=True)
registry_dict = fs.registry.to_dict(project=project)

# Save only the specs, not the metadata.
registry_specs = {
key: [fco["spec"] for fco in value]
for key, value in registry_dict.items()
}

# entity & feature view list commands should succeed
result = runner.run(["entities", "list"], cwd=repo_path)
Expand Down Expand Up @@ -83,8 +89,12 @@ def test_universal_cli(test_repo_config) -> None:
)

# Confirm that registry contents have not changed.
assertpy.assert_that(registry_dict).is_equal_to(
fs.registry.to_dict(project=project, spec_only=True)
registry_dict = fs.registry.to_dict(project=project)
assertpy.assert_that(registry_specs).is_equal_to(
{
key: [fco["spec"] for fco in value]
for key, value in registry_dict.items()
}
)

result = runner.run(["teardown"], cwd=repo_path)
Expand Down

0 comments on commit 4063d3d

Please sign in to comment.