def test_qc_data_set_collection_validation():
"""Check that the QC data set correctly validates for unique set ids."""
# Create a set which should not error
QCDataSetCollection(data_sets=[create_qc_data_set("qc-data-set-1")])
# Create a set which should error
with pytest.raises(ValidationError):
QCDataSetCollection(data_sets=[
create_qc_data_set("qc-data-set-1"),
create_qc_data_set("qc-data-set-1"),
])
def test_unique_entries(self, db: Session):
# Add data set with a single entry.
db.add(QCDataSetCRUD.create(db, create_qc_data_set("qc-data-set-1")))
db.commit()
assert db.query(models.QCDataSetEntry.id).count() == 1
# Make sure the previously added entry is referenced by the new data set.
db.add(QCDataSetCRUD.create(db, create_qc_data_set("qc-data-set-2")))
db.commit()
assert db.query(models.QCDataSetEntry.id).count() == 1
def test_qc_data_set_validation():
"""Check that the QC data set correctly validates for unique set ids."""
# Create a set which should not error
qc_data_set = create_qc_data_set("qc-data-set-1")
# Create a set which should error
with pytest.raises(ValidationError):
qc_data_set.entries.append(qc_data_set.entries[0])
QCDataSet(**qc_data_set.dict())
def test_find_or_retrieve_data_sets(
requests_mock, data_set_ids, data_set_type, expected_data_set_ids, expected_raises
):
local_data_sets = [
create_data_set("data-set-1", 1),
create_qc_data_set("data-set-1"),
create_data_set("data-set-2", 2),
]
remote_data_sets = [
create_qc_data_set("data-set-2"),
create_data_set("data-set-2", 3),
create_data_set("data-set-3", 4),
]
for remote_data_set in remote_data_sets:
if isinstance(remote_data_set, DataSet):
mock_get_data_set(requests_mock, remote_data_set)
else:
mock_get_qc_data_set(requests_mock, remote_data_set)
with expected_raises:
found_data_sets = InputFactory._find_or_retrieve_data_sets(
data_set_ids, data_set_type, local_data_sets
)
assert sorted(data_set_ids) == sorted(
data_set.id for data_set in found_data_sets
)
assert all(isinstance(data_set, data_set_type) for data_set in found_data_sets)
found_entry_ids = [
entry.id
for data_set in found_data_sets
if isinstance(data_set, DataSet)
for entry in data_set.entries
]
assert found_entry_ids == expected_data_set_ids
def commit_qc_data_set(db: Session,
unique_id: str = "qc-data-set-1") -> QCDataSet:
"""Creates a new QC data set and commits it the current session.
Parameters
----------
db
The current data base session.
unique_id
The id to assign to the QC data set.
"""
db_qc_data_set = QCDataSetCRUD.create(db, create_qc_data_set(unique_id))
db.add(db_qc_data_set)
db.commit()
return QCDataSetCRUD.db_to_model(db_qc_data_set)
def test_generate_recharge_target(self, requests_mock):
qc_data_set = create_qc_data_set("qc-data-set-1")
mock_get_qc_data_set(requests_mock, qc_data_set)
target = create_recharge_target("recharge-target-1", [qc_data_set.id])
with temporary_cd():
OptimizationInputFactory._generate_recharge_target(target, None)
with open("training-set.json") as file:
training_entries = json.load(file)
assert training_entries == qc_data_set.entries
with open("grid-settings.json") as file:
assert file.read() == target.grid_settings.json()
def _create_model(cls, db, create_dependencies=True):
qc_data_set = create_qc_data_set("qc-data-set-1")
return qc_data_set, {"qc_data_set_id": qc_data_set.id}
"mock-env",
"01:00",
"lilac-dask",
8000,
1,
False,
None,
optimization_result,
)
@pytest.mark.parametrize(
"reference_data_sets, expected_raises",
[
(
[create_data_set("data-set-1"), create_qc_data_set("data-set-1")],
does_not_raise(),
),
(
[
create_data_set("data-set-1"),
create_data_set("data-set-1"),
create_qc_data_set("data-set-1"),
],
pytest.raises(AssertionError, match="multiple reference data sets of"),
),
(
[
create_data_set("data-set-1"),
create_qc_data_set("data-set-1"),
create_qc_data_set("data-set-1"),
def create_dependencies(db: Session, dependencies: List[str]):
"""Create any dependencies such as parent studies, projects, or data sets and
commit them to the database.
Parameters
----------
db
The current database session.
dependencies
The required dependencies.
"""
project = None
data_set_ids = []
qc_data_set_ids = []
if "data-set" in dependencies:
data_set_ids.append("data-set-1")
if "qc-data-set" in dependencies:
qc_data_set_ids.append("qc-data-set-1")
for data_set_id in data_set_ids:
data_set = create_data_set(data_set_id)
db_data_set = DataSetCRUD.create(db, data_set)
db.add(db_data_set)
for qc_data_set_id in qc_data_set_ids:
qc_data_set = create_qc_data_set(qc_data_set_id)
db_qc_data_set = QCDataSetCRUD.create(db, qc_data_set)
db.add(db_qc_data_set)
db.commit()
if ("project" in dependencies or "study" in dependencies
or "evaluator-target" in dependencies
or "recharge-target" in dependencies
or "benchmark" in dependencies):
project = create_project("project-1")
if ("study" in dependencies or "evaluator-target" in dependencies
or "recharge-target" in dependencies
or "benchmark" in dependencies):
project.studies = [create_study(project.id, "study-1")]
if "evaluator-target" in dependencies or "recharge-target" in dependencies:
targets = []
if "evaluator-target" in dependencies:
targets.append(
create_evaluator_target("evaluator-target-1", ["data-set-1"]))
if "recharge-target" in dependencies:
targets.append(
create_recharge_target("recharge-target-1", ["qc-data-set-1"]))
optimization = create_optimization(project.id, project.studies[0].id,
"optimization-1", targets)
project.studies[0].optimizations = [optimization]
if "benchmark" in dependencies:
benchmark = create_benchmark(
project.id,
project.studies[0].id,
"benchmark-1",
["data-set-1"],
None,
create_force_field(),
)
project.studies[0].benchmarks = [benchmark]
if project is not None:
db_project = ProjectCRUD.create(db, project)
db.add(db_project)
db.commit()
def create_model(cls, include_children=False, index=1):
return create_qc_data_set(f"qc-data-set-{index}")