Fixing Unites functionality

state-manager/app/models/node_template_model.py

@@ -12,4 +12,4 @@ class NodeTemplate(BaseModel):
     identifier: str = Field(..., description="Identifier of the node")
     inputs: dict[str, Any] = Field(..., description="Inputs of the node")
     next_nodes: Optional[List[str]] = Field(None, description="Next nodes to execute")
-    unites: Optional[List[Unites]] = Field(None, description="Unites of the node")
+    unites: Optional[Unites] = Field(None, description="Unites of the node")

state-manager/app/tasks/create_next_states.py

@@ -40,24 +40,21 @@ async def mark_success_states(state_ids: list[PydanticObjectId]):
 
 
 async def check_unites_satisfied(namespace: str, graph_name: str, node_template: NodeTemplate, parents: dict[str, PydanticObjectId]) -> bool:
-    if node_template.unites is None or len(node_template.unites) == 0:
+    if node_template.unites is None:
         return True
 
-    for unit in node_template.unites:
-        unites_id = parents.get(unit.identifier)
-        if not unites_id:
-            raise ValueError(f"Unit identifier not found in parents: {unit.identifier}")
-        else:
-            pending_count = await State.find(
-                State.identifier == unit.identifier,
+    unites_id = parents.get(node_template.unites.identifier)
+    if not unites_id:
+        raise ValueError(f"Unit identifier not found in parents: {node_template.unites.identifier}")
+    else:
+        if await State.find(
                 State.namespace_name == namespace,
                 State.graph_name == graph_name,
                 NE(State.status, StateStatusEnum.SUCCESS),
                 {
-                    f"parents.{unit.identifier}": unites_id
+                    f"parents.{node_template.unites.identifier}": unites_id
                 }
-            ).count()
-            if pending_count > 0:
+            ).count() > 0:
                 return False
     return True
 
@@ -107,6 +104,41 @@ def validate_dependencies(next_state_node_template: NodeTemplate, next_state_inp
                     raise AttributeError(f"Output field '{dependent.field}' not found on state '{dependent.identifier}' for template '{next_state_node_template.identifier}'")
 
 
+def generate_next_state(next_state_input_model: Type[BaseModel], next_state_node_template: NodeTemplate, parents: dict[str, State], current_state: State) -> State:
+    next_state_input_data = {}
+
+    for field_name, _ in next_state_input_model.model_fields.items():
+        dependency_string = get_dependents(next_state_node_template.inputs[field_name])
+
+        for key in sorted(dependency_string.dependents.keys()):
+                if dependency_string.dependents[key].identifier == current_state.identifier:
+                    if dependency_string.dependents[key].field not in current_state.outputs:
+                        raise AttributeError(f"Output field '{dependency_string.dependents[key].field}' not found on current state '{current_state.identifier}' for template '{next_state_node_template.identifier}'")
+                    dependency_string.dependents[key].value = current_state.outputs[dependency_string.dependents[key].field]
+                else:
+                    dependency_string.dependents[key].value = parents[dependency_string.dependents[key].identifier].outputs[dependency_string.dependents[key].field]
+
+        next_state_input_data[field_name] = dependency_string.generate_string()
+
+    new_parents = {
+        **current_state.parents,
+        current_state.identifier: current_state.id
+    }
+
+    return State(
+        node_name=next_state_node_template.node_name,
+        identifier=next_state_node_template.identifier,
+        namespace_name=next_state_node_template.namespace,
+        graph_name=current_state.graph_name,
+        status=StateStatusEnum.CREATED,
+        parents=new_parents,
+        inputs=next_state_input_data,
+        outputs={},
+        run_id=current_state.run_id,
+        error=None
+    )
+
+
 async def create_next_states(state_ids: list[PydanticObjectId], identifier: str, namespace: str, graph_name: str, parents_ids: dict[str, PydanticObjectId]):
 
     try:
@@ -161,56 +193,47 @@ async def get_input_model(node_template: NodeTemplate) -> Type[BaseModel]:
         for parent_state in parent_states:
             parents[parent_state.identifier] = parent_state
 
+        pending_unites = []
 
         for next_state_identifier in next_state_identifiers:
             next_state_node_template = graph_template.get_node_by_identifier(next_state_identifier)
             if not next_state_node_template:
                 raise ValueError(f"Next state node template not found for identifier: {next_state_identifier}")
 
-            if not await check_unites_satisfied(namespace, graph_name, next_state_node_template, parents_ids):
+            if next_state_node_template.unites is not None:
+                pending_unites.append(next_state_identifier)
                 continue
 
             next_state_input_model = await get_input_model(next_state_node_template)
             validate_dependencies(next_state_node_template, next_state_input_model, identifier, parents)
 
             for current_state in current_states:                
-                next_state_input_data = {}
-
-                for field_name, _ in next_state_input_model.model_fields.items():
-                    dependency_string = get_dependents(next_state_node_template.inputs[field_name])
-
-                    for key in sorted(dependency_string.dependents.keys()):
-                        if dependency_string.dependents[key].identifier == identifier:
-                            if dependency_string.dependents[key].field not in current_state.outputs:
-                                raise AttributeError(f"Output field '{dependency_string.dependents[key].field}' not found on current state '{identifier}' for template '{next_state_node_template.identifier}'")
-                            dependency_string.dependents[key].value = current_state.outputs[dependency_string.dependents[key].field]
-                        else:
-                            dependency_string.dependents[key].value = parents[dependency_string.dependents[key].identifier].outputs[dependency_string.dependents[key].field]
-
-                    next_state_input_data[field_name] = dependency_string.generate_string()
-
-                new_parents = {
-                    **parents_ids,
-                    identifier: current_state.id
-                }
-
-                new_states.append(
-                    State(
-                        node_name=next_state_node_template.node_name,
-                        identifier=next_state_node_template.identifier,
-                        namespace_name=next_state_node_template.namespace,
-                        graph_name=graph_name,
-                        status=StateStatusEnum.CREATED,
-                        parents=new_parents,
-                        inputs=next_state_input_data,
-                        outputs={},
-                        run_id=current_state.run_id,
-                        error=None
-                    )
-                )
+                new_states.append(generate_next_state(next_state_input_model, next_state_node_template, parents, current_state))
 
-        await State.insert_many(new_states)
+        if len(new_states) > 0:
+            await State.insert_many(new_states)
         await mark_success_states(state_ids)
+
+        # handle unites
+        new_unit_states = []
+        for pending_unites_identifier in pending_unites:
+            next_state_node_template = graph_template.get_node_by_identifier(pending_unites_identifier)
+            if not next_state_node_template:
+                raise ValueError(f"Next state node template not found for identifier: {pending_unites_identifier}")
+
+            if not await check_unites_satisfied(namespace, graph_name, next_state_node_template, parents_ids):
+                continue
+
+            next_state_input_model = await get_input_model(next_state_node_template)
+            validate_dependencies(next_state_node_template, next_state_input_model, identifier, parents)
+
+            assert next_state_node_template.unites is not None
+            parent_state = parents[next_state_node_template.unites.identifier]
+
+            new_unit_states.append(generate_next_state(next_state_input_model, next_state_node_template, parents, parent_state))
+
+        if len(new_unit_states) > 0:
+            await State.insert_many(new_unit_states)
 
     except Exception as e:
         await State.find(

state-manager/app/tasks/verify_graph.py

@@ -224,11 +224,11 @@ async def verify_unites(graph_nodes: list[NodeTemplate], dependency_graph: dict
         return
 
     for node in graph_nodes:
-        if node.unites is None or len(node.unites) == 0:
+        if node.unites is None:
             continue
-        for depend in node.unites:
-            if depend.identifier not in dependency_graph[node.identifier]:
-                errors.append(f"Node {node.identifier} depends on {depend.identifier} which is not a dependency of {node.identifier}")
+
+        if node.unites.identifier not in dependency_graph[node.identifier]:
+            errors.append(f"Node {node.identifier} depends on {node.unites.identifier} which is not a dependency of {node.identifier}")
 
 
 async def verify_graph(graph_template: GraphTemplate):

state-manager/tests/unit/models/test_base.py

@@ -0,0 +1,55 @@
+import pytest
+from datetime import datetime
+from app.models.db.base import BaseDatabaseModel
+
+
+class TestBaseDatabaseModel:
+    """Test cases for BaseDatabaseModel"""
+
+    def test_base_model_field_definitions(self):
+        """Test that BaseDatabaseModel has the expected fields"""
+        # Check that the model has the expected fields
+        model_fields = BaseDatabaseModel.model_fields
+
+        assert 'created_at' in model_fields
+        assert 'updated_at' in model_fields
+
+        # Check field descriptions
+        assert model_fields['created_at'].description == "Date and time when the model was created"
+        assert model_fields['updated_at'].description == "Date and time when the model was last updated"
+
+    def test_base_model_abc_inheritance(self):
+        """Test that BaseDatabaseModel is an abstract base class"""
+        # Should not be able to instantiate BaseDatabaseModel directly
+        with pytest.raises(Exception):  # Could be TypeError or CollectionWasNotInitialized
+            BaseDatabaseModel()
+
+    def test_base_model_document_inheritance(self):
+        """Test that BaseDatabaseModel inherits from Document"""
+        # Check that it has the expected base classes
+        bases = BaseDatabaseModel.__bases__
+        assert len(bases) >= 2  # Should have at least ABC and Document as base classes
+
+    def test_base_model_has_update_updated_at_method(self):
+        """Test that BaseDatabaseModel has the update_updated_at method"""
+        assert hasattr(BaseDatabaseModel, 'update_updated_at')
+        assert callable(BaseDatabaseModel.update_updated_at)
+
+    def test_base_model_field_types(self):
+        """Test that BaseDatabaseModel fields have correct types"""
+        model_fields = BaseDatabaseModel.model_fields
+
+        # Check that created_at and updated_at are datetime fields
+        created_at_field = model_fields['created_at']
+        updated_at_field = model_fields['updated_at']
+
+        assert created_at_field.annotation == datetime
+        assert updated_at_field.annotation == datetime
+
+    def test_base_model_has_before_event_decorator(self):
+        """Test that BaseDatabaseModel uses the before_event decorator"""
+        # Check that the update_updated_at method exists and is callable
+        update_method = BaseDatabaseModel.update_updated_at
+
+        # The method should exist and be callable
+        assert callable(update_method) 

state-manager/tests/unit/models/test_graph_template_model.py

@@ -0,0 +1,107 @@
+import pytest
+from unittest.mock import patch
+import base64
+from app.models.db.graph_template_model import GraphTemplate
+
+
+class TestGraphTemplate:
+    """Test cases for GraphTemplate model"""
+
+    def test_validate_secrets_valid(self):
+        """Test validation of valid secrets"""
+        valid_secrets = {
+            "secret1": "valid_encrypted_string_that_is_long_enough_for_testing_32_chars",
+            "secret2": "another_valid_encrypted_string_that_is_long_enough_for_testing_32",
+        }
+
+        # Mock base64 decoding to succeed
+        with patch("base64.urlsafe_b64decode", return_value=b"x" * 20):
+            result = GraphTemplate.validate_secrets(valid_secrets)
+
+            assert result == valid_secrets
+
+    def test_validate_secrets_empty_name(self):
+        """Test validation with empty secret name"""
+        invalid_secrets = {"": "valid_value"}
+
+        with pytest.raises(ValueError, match="Secrets cannot be empty"):
+            GraphTemplate.validate_secrets(invalid_secrets)
+
+    def test_validate_secrets_empty_value(self):
+        """Test validation with empty secret value"""
+        invalid_secrets = {"secret1": ""}
+
+        with pytest.raises(ValueError, match="Secrets cannot be empty"):
+            GraphTemplate.validate_secrets(invalid_secrets)
+
+    def test_validate_secret_value_too_short(self):
+        """Test validation of secret value that's too short"""
+        short_value = "short"
+
+        with pytest.raises(ValueError, match="Value appears to be too short for an encrypted string"):
+            GraphTemplate._validate_secret_value(short_value)
+
+    def test_validate_secret_value_invalid_base64(self):
+        """Test validation of secret value with invalid base64"""
+        invalid_base64 = "invalid_base64_string_that_is_long_enough_but_not_valid_base64"
+
+        with pytest.raises(ValueError, match="Value is not valid URL-safe base64 encoded"):
+            GraphTemplate._validate_secret_value(invalid_base64)
+
+    def test_validate_secret_value_valid(self):
+        """Test validation of valid secret value"""
+        # Create a valid base64 string that decodes to at least 12 bytes and is long enough
+        valid_bytes = b"x" * 20
+        valid_base64 = base64.urlsafe_b64encode(valid_bytes).decode()
+        # Pad to make it at least 32 characters
+        padded_base64 = valid_base64 + "x" * (32 - len(valid_base64))
+
+        # Should not raise any exception
+        GraphTemplate._validate_secret_value(padded_base64)
+
+    def test_validate_secrets_with_long_valid_strings(self):
+        """Test validation with properly long secret values"""
+        long_secrets = {
+            "secret1": "x" * 50,  # 50 characters
+            "secret2": "y" * 100,  # 100 characters
+        }
+
+        # Mock base64 decoding to succeed
+        with patch("base64.urlsafe_b64decode", return_value=b"x" * 20):
+            result = GraphTemplate.validate_secrets(long_secrets)
+
+            assert result == long_secrets
+
+    def test_validate_secret_value_exactly_32_chars(self):
+        """Test validation with exactly 32 character string"""
+        exactly_32 = "x" * 32
+
+        # Mock base64 decoding to succeed
+        with patch("base64.urlsafe_b64decode", return_value=b"x" * 20):
+            # Should not raise exception
+            GraphTemplate._validate_secret_value(exactly_32)
+
+    def test_validate_secret_value_31_chars_fails(self):
+        """Test validation with 31 character string fails"""
+        exactly_31 = "x" * 31
+
+        with pytest.raises(ValueError, match="Value appears to be too short for an encrypted string"):
+            GraphTemplate._validate_secret_value(exactly_31)
+
+    def test_validate_secret_value_base64_decode_exception(self):
+        """Test validation when base64 decoding raises exception"""
+        invalid_base64 = "invalid_base64_string_that_is_long_enough_but_not_valid_base64"
+
+        with pytest.raises(ValueError, match="Value is not valid URL-safe base64 encoded"):
+            GraphTemplate._validate_secret_value(invalid_base64)
+
+    def test_validate_secret_value_decoded_exactly_12_bytes(self):
+        """Test validation with decoded value exactly 12 bytes"""
+        # Create a valid base64 string that decodes to exactly 12 bytes and is long enough
+        exactly_12_bytes = b"x" * 12
+        base64_string = base64.urlsafe_b64encode(exactly_12_bytes).decode()
+        # Pad to make it at least 32 characters
+        padded_base64 = base64_string + "x" * (32 - len(base64_string))
+
+        # Should not raise exception
+        GraphTemplate._validate_secret_value(padded_base64)