diff --git a/src/py/mat3ra/mode/methods/pseudopotential.py b/src/py/mat3ra/mode/methods/pseudopotential.py index e0d5a27..ffb6031 100644 --- a/src/py/mat3ra/mode/methods/pseudopotential.py +++ b/src/py/mat3ra/mode/methods/pseudopotential.py @@ -2,10 +2,11 @@ from mat3ra.esse.models.methods_directory.legacy.pseudopotential import LegacyMethodPseudopotential -from ..method import Method +from ..method import Method, MethodData class PseudopotentialMethod(LegacyMethodPseudopotential, Method): + data: MethodData = MethodData() pseudopotential_cls: Optional[type] = None @property @@ -32,8 +33,6 @@ def to_dict(self, exclude: Optional[List[str]] = None) -> Dict[str, Any]: json_data = super().to_dict(exclude=exclude) if exclude is None or "data" not in exclude: - filtered_data = self.data.copy() - # filtered_data.pop("allPseudo", None) - json_data["data"] = filtered_data + json_data["data"] = self.data.to_dict(exclude={"allPseudo"}) return json_data diff --git a/tests/py/unit/test_pseudopotential_method.py b/tests/py/unit/test_pseudopotential_method.py index 47b1185..8f46e43 100644 --- a/tests/py/unit/test_pseudopotential_method.py +++ b/tests/py/unit/test_pseudopotential_method.py @@ -40,13 +40,12 @@ def test_all_pseudo_property(config): assert isinstance(all_pseudo, list) assert len(all_pseudo) == 2 -@pytest.mark.skip(reason="TODO: fix this test") @pytest.mark.parametrize("config", TEST_CONFIGS) def test_to_dict_excludes_all_pseudo(config): config_with_data = {**config, "data": TEST_COMBINED_PSEUDO_DATA} method = PseudopotentialMethod.create(config_with_data) - json_data = method.model_dump() + json_data = method.to_dict() assert "allPseudo" not in json_data["data"] assert "pseudo" in json_data["data"]