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}")