diff --git a/tensormap-backend/app/services/data_process.py b/tensormap-backend/app/services/data_process.py index cab7df0b..f36b3c24 100644 --- a/tensormap-backend/app/services/data_process.py +++ b/tensormap-backend/app/services/data_process.py @@ -42,6 +42,14 @@ def add_target_service(db: Session, file_id: uuid_pkg.UUID, target: str) -> tupl file = db.exec(select(DataFile).where(DataFile.id == file_id)).first() if not file: return _resp(400, False, "File doesn't exist in DB") + if file.file_type != "csv": + return _resp(400, False, "Only CSV files support target field assignment") + if file.columns is not None and target not in file.columns: + return _resp( + 400, + False, + f"Target column '{target}' not found in dataset columns", + ) data_process = DataProcess(file_id=file_id, target=target) db.add(data_process) diff --git a/tensormap-backend/tests/test_data_process_service.py b/tensormap-backend/tests/test_data_process_service.py index 93074962..b8771e2c 100644 --- a/tensormap-backend/tests/test_data_process_service.py +++ b/tensormap-backend/tests/test_data_process_service.py @@ -110,6 +110,45 @@ def test_file_not_found(self, mock_db, file_id): assert body["success"] is False assert "doesn't exist" in body["message"] + def test_rejects_non_csv_file(self, mock_db, file_id): + zip_file = MagicMock(spec=DataFile) + zip_file.id = file_id + zip_file.file_type = "zip" + mock_db.exec.return_value.first.return_value = zip_file + + body, status = add_target_service(mock_db, file_id, "species") + + assert status == 400 + assert body["success"] is False + assert "CSV" in body["message"] + + def test_rejects_column_not_in_dataset(self, mock_db, file_id): + csv_file = MagicMock(spec=DataFile) + csv_file.id = file_id + csv_file.file_type = "csv" + csv_file.columns = ["sepal_length", "sepal_width", "species"] + mock_db.exec.return_value.first.return_value = csv_file + + body, status = add_target_service(mock_db, file_id, "nonexistent_column") + + assert status == 400 + assert body["success"] is False + assert "not found" in body["message"] + + def test_accepts_column_in_dataset(self, mock_db, file_id): + csv_file = MagicMock(spec=DataFile) + csv_file.id = file_id + csv_file.file_type = "csv" + csv_file.columns = ["sepal_length", "sepal_width", "species"] + mock_db.exec.return_value.first.return_value = csv_file + + body, status = add_target_service(mock_db, file_id, "species") + + assert status == 201 + assert body["success"] is True + mock_db.add.assert_called_once() + mock_db.commit.assert_called_once() + # --------------------------------------------------------------------------- # get_one_target_by_id_service