diff --git a/backend/data/base_model.py b/backend/data/base_model.py new file mode 100644 index 0000000..2ebc20d --- /dev/null +++ b/backend/data/base_model.py @@ -0,0 +1,16 @@ +from sqlmodel import SQLModel +from sqlmodel._compat import get_config_value, set_config_value + + +class BaseSQLModel(SQLModel): + """ + Base SQL Model class. It performs validation on the model unlike the default SQLModel class with table=True. + """ + + def __init__(self, **data): + is_table = get_config_value(model=self, parameter="table", default=False) + set_config_value( + model=self, parameter="table", value=False + ) # Makes it validate the model + super().__init__(**data) + set_config_value(model=self, parameter="table", value=is_table) diff --git a/backend/data/data_models.py b/backend/data/data_models.py index 6a9eab4..f5dfe01 100644 --- a/backend/data/data_models.py +++ b/backend/data/data_models.py @@ -2,12 +2,13 @@ # NOTE: This file should not be modified from datetime import datetime from pydantic import model_validator -from sqlmodel import Field, SQLModel +from sqlmodel import Field +from backend.data.base_model import BaseSQLModel from backend.data.enums import CommandStatus -class MainCommand(SQLModel, table=True): +class MainCommand(BaseSQLModel, table=True): """ Main command model. This table represents all the possible commands that can be issued. @@ -44,7 +45,7 @@ def validate_params_format(self): return self -class Command(SQLModel, table=True): +class Command(BaseSQLModel, table=True): """ An instance of a MainCommand. This table holds the data related to actual commands sent from the ground station up to the OBC. diff --git a/test/backend/test_data.py b/test/backend/test_data.py index e69de29..2295bd1 100644 --- a/test/backend/test_data.py +++ b/test/backend/test_data.py @@ -0,0 +1,57 @@ +import pytest +from sqlmodel import Session +from backend.data.data_models import MainCommand + + +def test_main_command(): + main_command = MainCommand( + name="Test", + params="param1,param2", + format="int,int", + data_size=2, + total_size=2, + ) + assert main_command.name == "Test" + assert main_command.params == "param1,param2" + assert main_command.format == "int,int" + assert main_command.data_size == 2 + assert main_command.total_size == 2 + + +def test_main_command_no_params_and_no_format(): + main_command = MainCommand( + name="Test", + data_size=2, + total_size=2, + ) + assert main_command.name == "Test" + assert main_command.params is None + assert main_command.format is None + assert main_command.data_size == 2 + assert main_command.total_size == 2 + + +def test_main_command_no_format(db_engine): + with pytest.raises(ValueError), Session(db_engine) as session: + session.add( + MainCommand( + name="Test", + params="param1,param2", + data_size=2, + total_size=2, + ) + ) + session.commit() + + +def test_main_command_no_params(db_engine): + with pytest.raises(ValueError), Session(db_engine) as session: + session.add( + MainCommand( + name="Test", + format="int,int", + data_size=2, + total_size=2, + ) + ) + session.commit()