diff --git a/gigl/src/validation_check/config_validator.py b/gigl/src/validation_check/config_validator.py index 49af4637f..ec0ca4caf 100644 --- a/gigl/src/validation_check/config_validator.py +++ b/gigl/src/validation_check/config_validator.py @@ -15,6 +15,10 @@ assert_subgraph_sampler_output_exists, assert_trained_model_exists, ) +from gigl.src.validation_check.libs.gbml_and_resource_config_compatibility_checks import ( + check_inferencer_graph_store_compatibility, + check_trainer_graph_store_compatibility, +) from gigl.src.validation_check.libs.name_checks import ( check_if_kfp_pipeline_job_name_valid, ) @@ -191,6 +195,79 @@ logger = Logger() +# Map of start components to graph store compatibility checks to run +# Only run trainer checks when starting at or before Trainer +# Only run inferencer checks when starting at or before Inferencer +START_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS = { + GiGLComponents.ConfigPopulator.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.DataPreprocessor.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.SubgraphSampler.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.SplitGenerator.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.Trainer.value: [ + check_trainer_graph_store_compatibility, + check_inferencer_graph_store_compatibility, + ], + GiGLComponents.Inferencer.value: [ + check_inferencer_graph_store_compatibility, + ], + # PostProcessor doesn't need graph store compatibility checks +} + +# Map of (start, stop) component tuples to graph store compatibility checks + +STOP_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS_TO_SKIP = { + GiGLComponents.Trainer.value: [ + check_inferencer_graph_store_compatibility, + ], +} + + +def _run_gbml_and_resource_config_compatibility_checks( + start_at: str, + stop_after: Optional[str], + gbml_config_pb_wrapper: GbmlConfigPbWrapper, + resource_config_wrapper: GiglResourceConfigWrapper, +) -> None: + """ + Run compatibility checks between GbmlConfig and GiglResourceConfig. + + These checks verify that graph store mode configurations are consistent + across both the template config (GbmlConfig) and resource config (GiglResourceConfig). + + Args: + start_at: The component to start at. + stop_after: Optional component to stop after. + gbml_config_pb_wrapper: The GbmlConfig wrapper (template config). + resource_config_wrapper: The GiglResourceConfig wrapper (resource config). + """ + # Get the appropriate compatibility checks based on start/stop components + compatibility_checks = set( + START_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS.get(start_at, []) + ) + if stop_after in STOP_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS_TO_SKIP: + for skipped_check in STOP_COMPONENT_TO_GRAPH_STORE_COMPATIBILITY_CHECKS_TO_SKIP[ + stop_after + ]: + compatibility_checks.discard(skipped_check) + + for check in compatibility_checks: + check( + gbml_config_pb_wrapper=gbml_config_pb_wrapper, + resource_config_wrapper=resource_config_wrapper, + ) + def kfp_validation_checks( job_name: str, @@ -261,6 +338,15 @@ def kfp_validation_checks( f"Skipping resource config check {resource_config_check.__name__} because we are using live subgraph sampling backend." ) + # check compatibility between template config and resource config for graph store mode + # These checks ensure that if graph store mode is enabled in one config, it's also enabled in the other + _run_gbml_and_resource_config_compatibility_checks( + start_at=start_at, + stop_after=stop_after, + gbml_config_pb_wrapper=gbml_config_pb_wrapper, + resource_config_wrapper=resource_config_wrapper, + ) + # check if trained model file exist when skipping training if gbml_config_pb.shared_config.should_skip_training == True: assert_trained_model_exists(gbml_config_pb=gbml_config_pb) diff --git a/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py new file mode 100644 index 000000000..fc12d1939 --- /dev/null +++ b/gigl/src/validation_check/libs/gbml_and_resource_config_compatibility_checks.py @@ -0,0 +1,137 @@ +""" +Compatibility checks between GbmlConfig (template config) and GiglResourceConfig (resource config). + +These checks ensure that graph store mode configurations are consistent across both configs. +If graph store mode is set up for trainer or inferencer in one config, it must be set up in the other. +""" + +from typing import Literal + +from google.protobuf.message import Message + +from gigl.common.logger import Logger +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper +from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( + GiglResourceConfigWrapper, +) +from snapchat.research.gbml import gigl_resource_config_pb2 + +logger = Logger() + + +def _gbml_config_has_graph_store( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, + component: Literal["trainer", "inferencer"], +) -> bool: + """ + Check if the GbmlConfig has graph_store_storage_config set for inferencer. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper to check. + + Returns: + True if graph_store_storage_config is set for inferencer, False otherwise. + """ + if component == "inferencer": + config: Message = gbml_config_pb_wrapper.gbml_config_pb.inferencer_config + elif component == "trainer": + config = gbml_config_pb_wrapper.gbml_config_pb.trainer_config + else: + raise ValueError( + f"Invalid component: {component}. Must be 'inferencer' or 'trainer'." + ) + return config.HasField("graph_store_storage_config") + + +def _resource_config_has_graph_store( + resource_config_wrapper: GiglResourceConfigWrapper, + component: Literal["trainer", "inferencer"], +) -> bool: + """ + Check if the GiglResourceConfig has VertexAiGraphStoreConfig set for the given component. + + Args: + resource_config_wrapper: The resource config wrapper to check. + + Returns: + True if VertexAiGraphStoreConfig is set for trainer, False otherwise. + """ + if component == "trainer": + config: Message = resource_config_wrapper.trainer_config + elif component == "inferencer": + config = resource_config_wrapper.inferencer_config + else: + raise ValueError( + f"Invalid component: {component}. Must be 'trainer' or 'inferencer'." + ) + return isinstance(config, gigl_resource_config_pb2.VertexAiGraphStoreConfig) + + +def check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, + resource_config_wrapper: GiglResourceConfigWrapper, +) -> None: + """ + Check that trainer graph store mode is consistently configured across both configs. + + If graph_store_storage_config is set in GbmlConfig.trainer_config, then + VertexAiGraphStoreConfig must be set in GiglResourceConfig.trainer_resource_config, + and vice versa. Also validates that storage_command is set when graph store mode is enabled. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper (template config). + resource_config_wrapper: The GiglResourceConfig wrapper (resource config). + + Raises: + AssertionError: If graph store configurations are not compatible or storage_command is missing. + """ + logger.info( + "Config validation check: trainer graph store compatibility between template and resource configs." + ) + + gbml_has_graph_store = _gbml_config_has_graph_store( + gbml_config_pb_wrapper, "trainer" + ) + resource_has_graph_store = _resource_config_has_graph_store( + resource_config_wrapper, "trainer" + ) + + if gbml_has_graph_store != resource_has_graph_store: + raise AssertionError( + f"If one of GbmlConfig.trainer_config.graph_store_storage_config or GiglResourceConfig.trainer_resource_config is set, the other must also be set. GbmlConfig.trainer_config.graph_store_storage_config is set: {gbml_has_graph_store}, GiglResourceConfig.trainer_resource_config is set: {resource_has_graph_store}." + ) + + +def check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, + resource_config_wrapper: GiglResourceConfigWrapper, +) -> None: + """ + Check that inferencer graph store mode is consistently configured across both configs. + + If graph_store_storage_config is set in GbmlConfig.inferencer_config, then + VertexAiGraphStoreConfig must be set in GiglResourceConfig.inferencer_resource_config, + and vice versa. Also validates that storage_command is set when graph store mode is enabled. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper (template config). + resource_config_wrapper: The GiglResourceConfig wrapper (resource config). + + Raises: + AssertionError: If graph store configurations are not compatible or storage_command is missing. + """ + logger.info( + "Config validation check: inferencer graph store compatibility between template and resource configs." + ) + + gbml_has_graph_store = _gbml_config_has_graph_store( + gbml_config_pb_wrapper, "inferencer" + ) + resource_has_graph_store = _resource_config_has_graph_store( + resource_config_wrapper, "inferencer" + ) + + if gbml_has_graph_store != resource_has_graph_store: + raise AssertionError( + f"If one of GbmlConfig.inferencer_config.graph_store_storage_config or GiglResourceConfig.inferencer_resource_config is set, the other must also be set. GbmlConfig.inferencer_config.graph_store_storage_config is set: {gbml_has_graph_store}, GiglResourceConfig.inferencer_resource_config is set: {resource_has_graph_store}." + ) diff --git a/gigl/src/validation_check/libs/resource_config_checks.py b/gigl/src/validation_check/libs/resource_config_checks.py index a60a66945..8e1137726 100644 --- a/gigl/src/validation_check/libs/resource_config_checks.py +++ b/gigl/src/validation_check/libs/resource_config_checks.py @@ -3,6 +3,7 @@ from google.cloud.aiplatform_v1.types.accelerator_type import AcceleratorType from gigl.common.logger import Logger +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( GiglResourceConfigWrapper, ) @@ -254,3 +255,53 @@ def _validate_machine_config( or {gigl_resource_config_pb2.VertexAiGraphStoreConfig.__name__}. Got {type(config)}""" ) + + +def check_if_trainer_graph_store_storage_command_valid( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, +) -> None: + """ + Validates that storage_command is set when graph store mode is enabled for trainer. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper to check. + + Raises: + AssertionError: If graph store mode is enabled but storage_command is missing. + """ + logger.info( + "Config validation check: if trainer graph store storage_command is valid." + ) + trainer_config = gbml_config_pb_wrapper.gbml_config_pb.trainer_config + if trainer_config.HasField("graph_store_storage_config"): + storage_command = trainer_config.graph_store_storage_config.command + if not storage_command: + raise AssertionError( + "GbmlConfig.trainer_config.graph_store_storage_config.storage_command must be set " + "when using graph store mode for trainer." + ) + + +def check_if_inferencer_graph_store_storage_command_valid( + gbml_config_pb_wrapper: GbmlConfigPbWrapper, +) -> None: + """ + Validates that storage_command is set when graph store mode is enabled for inferencer. + + Args: + gbml_config_pb_wrapper: The GbmlConfig wrapper to check. + + Raises: + AssertionError: If graph store mode is enabled but storage_command is missing. + """ + logger.info( + "Config validation check: if inferencer graph store storage_command is valid." + ) + inferencer_config = gbml_config_pb_wrapper.gbml_config_pb.inferencer_config + if inferencer_config.HasField("graph_store_storage_config"): + storage_command = inferencer_config.graph_store_storage_config.command + if not storage_command: + raise AssertionError( + "GbmlConfig.inferencer_config.graph_store_storage_config.storage_command must be set " + "when using graph store mode for inferencer." + ) diff --git a/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py b/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py new file mode 100644 index 000000000..49966df63 --- /dev/null +++ b/tests/unit/src/validation/lib/gbml_and_resource_config_compatibility_checks_test.py @@ -0,0 +1,217 @@ +import unittest + +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper +from gigl.src.common.types.pb_wrappers.gigl_resource_config import ( + GiglResourceConfigWrapper, +) +from gigl.src.validation_check.libs.gbml_and_resource_config_compatibility_checks import ( + check_inferencer_graph_store_compatibility, + check_trainer_graph_store_compatibility, +) +from snapchat.research.gbml import gbml_config_pb2, gigl_resource_config_pb2 + +# Helper functions for creating VertexAiGraphStoreConfig + + +def _create_vertex_ai_graph_store_config() -> ( + gigl_resource_config_pb2.VertexAiGraphStoreConfig +): + """Create a valid VertexAiGraphStoreConfig.""" + config = gigl_resource_config_pb2.VertexAiGraphStoreConfig() + # Graph store pool + config.graph_store_pool.machine_type = "n1-highmem-8" + config.graph_store_pool.gpu_type = "ACCELERATOR_TYPE_UNSPECIFIED" + config.graph_store_pool.gpu_limit = 0 + config.graph_store_pool.num_replicas = 2 + # Compute pool + config.compute_pool.machine_type = "n1-standard-16" + config.compute_pool.gpu_type = "NVIDIA_TESLA_T4" + config.compute_pool.gpu_limit = 2 + config.compute_pool.num_replicas = 3 + return config + + +def _create_vertex_ai_resource_config() -> ( + gigl_resource_config_pb2.VertexAiResourceConfig +): + """Create a valid VertexAiResourceConfig (non-graph store).""" + config = gigl_resource_config_pb2.VertexAiResourceConfig() + config.machine_type = "n1-standard-16" + config.gpu_type = "NVIDIA_TESLA_T4" + config.gpu_limit = 2 + config.num_replicas = 3 + return config + + +def _create_shared_resource_config( + config: gigl_resource_config_pb2.GiglResourceConfig, +) -> None: + """Populate shared resource config fields.""" + config.shared_resource_config.common_compute_config.project = "test-project" + config.shared_resource_config.common_compute_config.region = "us-central1" + config.shared_resource_config.common_compute_config.temp_assets_bucket = ( + "gs://test-temp" + ) + config.shared_resource_config.common_compute_config.temp_regional_assets_bucket = ( + "gs://test-temp-regional" + ) + config.shared_resource_config.common_compute_config.perm_assets_bucket = ( + "gs://test-perm" + ) + config.shared_resource_config.common_compute_config.temp_assets_bq_dataset_name = ( + "test_dataset" + ) + config.shared_resource_config.common_compute_config.embedding_bq_dataset_name = ( + "test_embeddings" + ) + config.shared_resource_config.common_compute_config.gcp_service_account_email = ( + "test@test-project.iam.gserviceaccount.com" + ) + config.shared_resource_config.common_compute_config.dataflow_runner = ( + "DataflowRunner" + ) + + +# Helper functions for creating GbmlConfig configurations + + +def _create_gbml_config_with_both_graph_stores( + storage_command: str = "python -m gigl.distributed.graph_store.storage_main", +) -> GbmlConfigPbWrapper: + """Create a GbmlConfig with graph_store_storage_config set for both trainer and inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.graph_store_storage_config.command = storage_command + gbml_config.inferencer_config.graph_store_storage_config.command = storage_command + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_without_graph_stores() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for trainer or inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" + gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_resource_config_with_both_graph_stores() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig with VertexAiGraphStoreConfig for both trainer and inferencer.""" + config = gigl_resource_config_pb2.GiglResourceConfig() + _create_shared_resource_config(config) + + # Trainer with VertexAiGraphStoreConfig + config.trainer_resource_config.vertex_ai_graph_store_trainer_config.CopyFrom( + _create_vertex_ai_graph_store_config() + ) + # Inferencer with VertexAiGraphStoreConfig + config.inferencer_resource_config.vertex_ai_graph_store_inferencer_config.CopyFrom( + _create_vertex_ai_graph_store_config() + ) + return GiglResourceConfigWrapper(resource_config=config) + + +def _create_resource_config_without_graph_stores() -> GiglResourceConfigWrapper: + """Create a GiglResourceConfig without VertexAiGraphStoreConfig for trainer or inferencer.""" + config = gigl_resource_config_pb2.GiglResourceConfig() + _create_shared_resource_config(config) + + # Trainer with standard config + config.trainer_resource_config.vertex_ai_trainer_config.CopyFrom( + _create_vertex_ai_resource_config() + ) + # Inferencer with standard config + config.inferencer_resource_config.vertex_ai_inferencer_config.CopyFrom( + _create_vertex_ai_resource_config() + ) + return GiglResourceConfigWrapper(resource_config=config) + + +class TestTrainerGraphStoreCompatibility(unittest.TestCase): + """Test suite for trainer graph store compatibility checks.""" + + def test_both_have_trainer_graph_store(self): + """Test that both configs having trainer graph store passes validation.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() + # Should not raise any exception + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_neither_has_trainer_graph_store(self): + """Test that neither config having trainer graph store passes validation.""" + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_without_graph_stores() + # Should not raise any exception + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_template_has_trainer_graph_store_resource_does_not(self): + """Test that template having graph store but resource not raises an assertion error.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_without_graph_stores() + with self.assertRaises(AssertionError): + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_resource_has_trainer_graph_store_template_does_not(self): + """Test that resource having graph store but template not raises an assertion error.""" + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() + with self.assertRaises(AssertionError): + check_trainer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + +class TestInferencerGraphStoreCompatibility(unittest.TestCase): + """Test suite for inferencer graph store compatibility checks.""" + + def test_both_have_inferencer_graph_store(self): + """Test that both configs having inferencer graph store passes validation.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() + # Should not raise any exception + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_neither_has_inferencer_graph_store(self): + """Test that neither config having inferencer graph store passes validation.""" + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_without_graph_stores() + # Should not raise any exception + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_template_has_inferencer_graph_store_resource_does_not(self): + """Test that template having graph store but resource not raises an assertion error.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + resource_config = _create_resource_config_without_graph_stores() + with self.assertRaises(AssertionError): + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + def test_resource_has_inferencer_graph_store_template_does_not(self): + """Test that resource having graph store but template not raises an assertion error.""" + gbml_config = _create_gbml_config_without_graph_stores() + resource_config = _create_resource_config_with_both_graph_stores() + with self.assertRaises(AssertionError): + check_inferencer_graph_store_compatibility( + gbml_config_pb_wrapper=gbml_config, + resource_config_wrapper=resource_config, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/unit/src/validation/lib/resource_config_checks_test.py b/tests/unit/src/validation/lib/resource_config_checks_test.py index 20b941f72..d136b85a8 100644 --- a/tests/unit/src/validation/lib/resource_config_checks_test.py +++ b/tests/unit/src/validation/lib/resource_config_checks_test.py @@ -1,18 +1,21 @@ import unittest +from gigl.src.common.types.pb_wrappers.gbml_config import GbmlConfigPbWrapper from gigl.src.validation_check.libs.resource_config_checks import ( _check_if_dataflow_resource_config_valid, _check_if_spark_resource_config_valid, _validate_accelerator_type, _validate_machine_config, + check_if_inferencer_graph_store_storage_command_valid, check_if_inferencer_resource_config_valid, check_if_preprocessor_resource_config_valid, check_if_shared_resource_config_valid, check_if_split_generator_resource_config_valid, check_if_subgraph_sampler_resource_config_valid, + check_if_trainer_graph_store_storage_command_valid, check_if_trainer_resource_config_valid, ) -from snapchat.research.gbml import gigl_resource_config_pb2 +from snapchat.research.gbml import gbml_config_pb2, gigl_resource_config_pb2 # Helper functions for creating valid configurations @@ -750,5 +753,70 @@ def test_valid_vertex_ai_graph_store_config(self): _validate_machine_config(config) +# Helper functions for creating GbmlConfig configurations + + +def _create_gbml_config_with_both_graph_stores( + storage_command: str = "python -m gigl.distributed.graph_store.storage_main", +) -> GbmlConfigPbWrapper: + """Create a GbmlConfig with graph_store_storage_config set for both trainer and inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.graph_store_storage_config.command = storage_command + gbml_config.inferencer_config.graph_store_storage_config.command = storage_command + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +def _create_gbml_config_without_graph_stores() -> GbmlConfigPbWrapper: + """Create a GbmlConfig without graph_store_storage_config for trainer or inferencer.""" + gbml_config = gbml_config_pb2.GbmlConfig() + gbml_config.trainer_config.trainer_args["some_arg"] = "some_value" + gbml_config.inferencer_config.inferencer_args["some_arg"] = "some_value" + return GbmlConfigPbWrapper(gbml_config_pb=gbml_config) + + +class TestTrainerGraphStoreStorageCommand(unittest.TestCase): + """Test suite for trainer graph store storage_command validation.""" + + def test_valid_storage_command(self): + """Test that a valid storage_command passes validation.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + # Should not raise any exception + check_if_trainer_graph_store_storage_command_valid(gbml_config) + + def test_missing_storage_command(self): + """Test that missing storage_command raises an assertion error.""" + gbml_config = _create_gbml_config_with_both_graph_stores(storage_command="") + with self.assertRaises(AssertionError): + check_if_trainer_graph_store_storage_command_valid(gbml_config) + + def test_no_graph_store_config(self): + """Test that no graph store config passes validation (nothing to check).""" + gbml_config = _create_gbml_config_without_graph_stores() + # Should not raise any exception - no graph store means nothing to validate + check_if_trainer_graph_store_storage_command_valid(gbml_config) + + +class TestInferencerGraphStoreStorageCommand(unittest.TestCase): + """Test suite for inferencer graph store storage_command validation.""" + + def test_valid_storage_command(self): + """Test that a valid storage_command passes validation.""" + gbml_config = _create_gbml_config_with_both_graph_stores() + # Should not raise any exception + check_if_inferencer_graph_store_storage_command_valid(gbml_config) + + def test_missing_storage_command(self): + """Test that missing storage_command raises an assertion error.""" + gbml_config = _create_gbml_config_with_both_graph_stores(storage_command="") + with self.assertRaises(AssertionError): + check_if_inferencer_graph_store_storage_command_valid(gbml_config) + + def test_no_graph_store_config(self): + """Test that no graph store config passes validation (nothing to check).""" + gbml_config = _create_gbml_config_without_graph_stores() + # Should not raise any exception - no graph store means nothing to validate + check_if_inferencer_graph_store_storage_command_valid(gbml_config) + + if __name__ == "__main__": unittest.main()