1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153
|
import os
from core_function import client_search_testing
import pytest
from pymatgen.analysis.magnetism import Ordering
from mp_api.client.routes.materials.electronic_structure import (
BandStructureRester,
DosRester,
ElectronicStructureRester,
)
@pytest.fixture
def es_rester():
rester = ElectronicStructureRester()
yield rester
rester.session.close()
es_excluded_params = [
"sort_fields",
"chunk_size",
"num_chunks",
"all_fields",
"fields",
]
sub_doc_fields = [] # type: list
es_alt_name_dict = {
"material_ids": "material_id",
"exclude_elements": "material_id",
"formula": "material_id",
"num_elements": "nelements",
"num_sites": "nsites",
} # type: dict
es_custom_field_tests = {
"material_ids": ["mp-149"],
"magnetic_ordering": Ordering.FM,
"formula": "CoO2",
"chemsys": "Co-O",
"elements": ["Co", "O"],
"exclude_elements": ["Co"],
} # type: dict
@pytest.mark.skipif(os.getenv("MP_API_KEY", None) is None, reason="No API key found.")
@pytest.mark.skip(reason="magnetic ordering fields not build correctly")
def test_es_client(es_rester):
search_method = es_rester.search
client_search_testing(
search_method=search_method,
excluded_params=es_excluded_params,
alt_name_dict=es_alt_name_dict,
custom_field_tests=es_custom_field_tests,
sub_doc_fields=sub_doc_fields,
)
bs_custom_field_tests = {
"magnetic_ordering": Ordering.FM,
"is_metal": True,
"is_gap_direct": True,
"efermi": (0, 100),
"band_gap": (0, 5),
}
bs_sub_doc_fields = ["bandstructure"]
bs_alt_name_dict = {} # type: dict
@pytest.fixture
def bs_rester():
rester = BandStructureRester()
yield rester
rester.session.close()
@pytest.mark.skipif(os.getenv("MP_API_KEY", None) is None, reason="No API key found.")
@pytest.mark.skip(reason="magnetic ordering fields not build correctly")
def test_bs_client(bs_rester):
# Get specific search method
search_method = bs_rester.search
# Query fields
for param in bs_custom_field_tests:
project_field = bs_alt_name_dict.get(param, None)
q = {
param: bs_custom_field_tests[param],
"chunk_size": 1,
"num_chunks": 1,
}
doc = search_method(**q)[0].model_dump()
for sub_field in bs_sub_doc_fields:
if sub_field in doc:
doc = doc[sub_field]
if param != "path_type":
doc = doc["setyawan_curtarolo"]
assert doc[project_field if project_field is not None else param] is not None
dos_custom_field_tests = {
"magnetic_ordering": Ordering.FM,
"efermi": (0, 100),
"band_gap": (0, 5),
}
dos_excluded_params = ["orbital", "element"]
dos_sub_doc_fields = ["dos"]
dos_alt_name_dict = {} # type: dict
@pytest.fixture
def dos_rester():
rester = DosRester()
yield rester
rester.session.close()
@pytest.mark.skipif(os.getenv("MP_API_KEY", None) is None, reason="No API key found.")
@pytest.mark.skip(reason="magnetic ordering fields not build correctly")
def test_dos_client(dos_rester):
search_method = dos_rester.search
# Query fields
for param in dos_custom_field_tests:
if param not in dos_excluded_params:
project_field = dos_alt_name_dict.get(param, None)
q = {
param: dos_custom_field_tests[param],
"chunk_size": 1,
"num_chunks": 1,
}
doc = search_method(**q)[0].model_dump()
for sub_field in dos_sub_doc_fields:
if sub_field in doc:
doc = doc[sub_field]
if param != "projection_type" and param != "magnetic_ordering":
doc = doc["total"]["1"]
assert (
doc[project_field if project_field is not None else param] is not None
)
|