github-starred/flowsint
2
0
Fork 0
mirror of https://github.com/reconurge/flowsint.git synced 2026-07-20 20:43:44 -05:00
Code Issues 980 Packages Projects Releases 14 Wiki Activity

feat(core): only support pydantic object insert and manipulation

Browse Source
This commit is contained in:
dextmorgn
2025-12-05 15:45:23 +01:00
parent 8d74c11506
commit 7562da12a6
11 changed files with 914 additions and 809 deletions
Download Patch File Download Diff File Expand all files Collapse all files
flowsint-core/src/flowsint_core/core/enricher_base.py
+16 -207
View File
@@ -8,6 +8,7 @@ from .vault import VaultProtocol
from .graph_service import GraphService, create_graph_service
from ..utils import resolve_type
class InvalidEnricherParams(Exception):
pass
@@ -413,9 +414,7 @@ class Enricher(ABC):
async def execute(self, values: List[str]) -> List[Dict[str, Any]]:
if self.name() != "enricher_orchestrator":
Logger.info(
self.sketch_id, {"message": f"Enricher {self.name()} started."}
)
Logger.info(self.sketch_id, {"message": f"Enricher {self.name()} started."})
try:
await self.async_init()
preprocessed = self.preprocess(values)
@@ -440,240 +439,50 @@ class Enricher(ABC):
)
return []
def create_node(
self, node_type_or_obj, key_prop=None, key_value=None, **properties
) -> None:
def create_node(self, node_obj, **properties) -> None:
"""
Create a single Neo4j node.
This method now uses the GraphService for improved performance and
better separation of concerns.
The following properties are automatically added to every node:
- type: Lowercase version of node_type
- sketch_id: Current sketch ID from enricher context
- label: Automatically computed by FlowsintType, or defaults to key_value if not provided
- created_at: ISO 8601 UTC timestamp (only on creation, not updates)
Best Practice - Use Pydantic object directly:
The simplest way is to pass a Pydantic object directly. The node type,
key property, and key value are automatically inferred:
Use Pydantic object directly:
```python
# Best: pass the Pydantic object directly
self.create_node(ip)
# Also good if you need to override properties
self.create_node(domain, type="subdomain")
```
Args:
node_type_or_obj: Either a Pydantic object (FlowsintType), or node label string (e.g., "domain", "ip")
key_prop: Property name used as unique identifier (optional if passing Pydantic object)
key_value: Value of the key property (optional if passing Pydantic object)
**properties: Additional node properties or property overrides
Note:
Uses MERGE semantics - if a node with the same (key_prop, sketch_id)
exists, it will be updated. The created_at field is only set on creation.
node_obj: Either a Pydantic object or node label string
**properties: Additional node properties or overrides
"""
# Check if first argument is a Pydantic object
if isinstance(node_type_or_obj, BaseModel):
obj = node_type_or_obj
# Infer node_type from class name (e.g., Ip -> "ip", Domain -> "domain")
node_type = obj.__class__.__name__.lower()
# Get the primary field and its value
primary_field = self._get_primary_field(obj)
key_prop = primary_field
key_value = getattr(obj, primary_field)
# If key_value is itself a Pydantic model, extract its primary value
if isinstance(key_value, BaseModel):
key_value = self._extract_primary_value(key_value)
# Merge object properties with any overrides, but skip nested Pydantic objects
# Use model_dump(mode="json") to properly serialize Pydantic types (e.g., HttpUrl)
obj_dict = obj.model_dump(mode="json") if hasattr(obj, "model_dump") else obj.dict()
obj_properties = {}
for k, v in obj_dict.items():
# Skip nested Pydantic objects (represented as dicts after model_dump)
if not isinstance(v, dict):
obj_properties[k] = v
obj_properties.update(properties)
properties = obj_properties
else:
# Legacy signature: node_type_or_obj is the node_type string
node_type = node_type_or_obj
self._graph_service.create_node(
node_type=node_type, key_prop=key_prop, key_value=key_value, **properties
)
def _serialize_properties(self, properties: dict) -> dict:
"""
Convert properties to Neo4j-compatible values.
DEPRECATED: This method is kept for backward compatibility.
New code should use GraphSerializer directly.
Args:
properties: Dictionary of properties to serialize
Returns:
Dictionary of serialized properties
"""
from .graph_serializer import GraphSerializer
return GraphSerializer.serialize_properties(properties)
self._graph_service.create_node(node_obj=node_obj, **properties)
def create_relationship(
self,
from_type_or_obj,
from_key_or_to_obj,
from_value_or_rel_type=None,
to_type=None,
to_key=None,
to_value=None,
rel_type=None,
from_obj,
to_obj,
rel_label="IS_RELATED_TO",
) -> None:
"""
Create a relationship between two nodes.
This method now uses the GraphService for improved performance and
better separation of concerns.
Best Practice - Use Pydantic objects directly:
The simplest way is to pass two Pydantic objects and the relationship type:
```python
# Best: pass Pydantic objects directly
self.create_relationship(individual, domain, "HAS_DOMAIN")
self.create_relationship(email, breach, "FOUND_IN_BREACH")
```
Legacy Usage:
You can still use the explicit signature for backward compatibility:
```python
# Legacy: explicit signature
self.create_relationship(
"individual", "full_name", individual.full_name,
"domain", "domain", domain_name,
"HAS_DOMAIN"
)
```
Args:
from_type_or_obj: Either a Pydantic object (source node) or source node label string
from_key_or_to_obj: Either a Pydantic object (target node) or source node key property
from_value_or_rel_type: Either relationship type string (if using objects) or source node key value
to_type: Target node label (only for legacy signature)
to_key: Target node key property (only for legacy signature)
to_value: Target node key value (only for legacy signature)
rel_type: Relationship type (only for legacy signature)
from_obj: Either a Pydantic object (source) or source node label
to_obj: Either a Pydantic object (target) or source node key property
rel_label: Either relationship type (Pydantic) or source node key value
"""
# Check if using new signature (Pydantic objects)
if isinstance(from_type_or_obj, BaseModel) and isinstance(from_key_or_to_obj, BaseModel):
from_obj = from_type_or_obj
to_obj = from_key_or_to_obj
relationship_type = from_value_or_rel_type
# Extract from_node info
from_node_type = from_obj.__class__.__name__.lower()
from_primary_field = self._get_primary_field(from_obj)
# Use model_dump to properly serialize Pydantic types (e.g., HttpUrl)
from_obj_dict = from_obj.model_dump(mode="json") if hasattr(from_obj, "model_dump") else from_obj.dict()
from_key_value = from_obj_dict.get(from_primary_field)
# If key_value is still a dict (nested Pydantic model), extract its primary value
if isinstance(from_key_value, dict):
# Get the raw nested object to extract its primary value
nested_obj = getattr(from_obj, from_primary_field)
if isinstance(nested_obj, BaseModel):
from_key_value = self._extract_primary_value(nested_obj)
# Extract to_node info
to_node_type = to_obj.__class__.__name__.lower()
to_primary_field = self._get_primary_field(to_obj)
# Use model_dump to properly serialize Pydantic types (e.g., HttpUrl)
to_obj_dict = to_obj.model_dump(mode="json") if hasattr(to_obj, "model_dump") else to_obj.dict()
to_key_value = to_obj_dict.get(to_primary_field)
# If key_value is still a dict (nested Pydantic model), extract its primary value
if isinstance(to_key_value, dict):
# Get the raw nested object to extract its primary value
nested_obj = getattr(to_obj, to_primary_field)
if isinstance(nested_obj, BaseModel):
to_key_value = self._extract_primary_value(nested_obj)
self._graph_service.create_relationship(
from_type=from_node_type,
from_key=from_primary_field,
from_value=from_key_value,
to_type=to_node_type,
to_key=to_primary_field,
to_value=to_key_value,
rel_type=relationship_type,
)
else:
# Legacy signature
self._graph_service.create_relationship(
from_type=from_type_or_obj,
from_key=from_key_or_to_obj,
from_value=from_value_or_rel_type,
to_type=to_type,
to_key=to_key,
to_value=to_value,
rel_type=rel_type,
)
def _get_primary_field(self, obj: BaseModel) -> str:
"""Helper method to get the primary field of a Pydantic object."""
# Access model_fields from the class, not the instance
model_fields = obj.__class__.model_fields
# Find the primary field (marked with json_schema_extra={"primary": True})
primary_field = None
for field_name, field_info in model_fields.items():
if field_info.json_schema_extra and field_info.json_schema_extra.get("primary"):
primary_field = field_name
break
# Fallback: use first required field or first field
if primary_field is None:
for field_name, field_info in model_fields.items():
if field_info.is_required():
primary_field = field_name
break
if primary_field is None:
primary_field = next(iter(model_fields.keys()))
return primary_field
def _extract_primary_value(self, obj: BaseModel) -> Any:
"""
Extract the primitive value from a Pydantic object recursively.
If the primary field is itself a Pydantic object, recursively extract its primary value.
Uses model_dump to properly serialize Pydantic types like HttpUrl.
"""
primary_field = self._get_primary_field(obj)
# Use model_dump to properly serialize Pydantic types (e.g., HttpUrl)
obj_dict = obj.model_dump(mode="json") if hasattr(obj, "model_dump") else obj.dict()
value = obj_dict.get(primary_field)
# If the value is still a dict (nested Pydantic model), recursively extract
if isinstance(value, dict):
# Get the raw nested object to recursively extract
nested_obj = getattr(obj, primary_field)
if isinstance(nested_obj, BaseModel):
return self._extract_primary_value(nested_obj)
return value
self._graph_service.create_relationship(
from_obj=from_obj, to_obj=to_obj, rel_label=rel_label
)
def log_graph_message(self, message: str) -> None:
"""
flowsint-core/src/flowsint_core/core/graph_db.py
+13 -3
View File
@@ -91,9 +91,17 @@ class Neo4jConnection:
List of result records as dictionaries
"""
with self._driver.session() as session:
result = session.run(query, parameters or {})
cleaned_params = self._clean_parameters(parameters)
result = session.run(query, cleaned_params)
return result.data()
@staticmethod
def _clean_parameters(parameters: Dict[str, Any]) -> Dict[str, Any]:
"""Remove None keys from parameters dict to avoid Neo4j errors."""
if not parameters:
return {}
return {k: v for k, v in parameters.items() if k is not None}
def execute_write(self, query: str, parameters: Dict[str, Any] = None) -> List[Dict[str, Any]]:
"""
Execute a write query within a write transaction.
@@ -106,7 +114,8 @@ class Neo4jConnection:
List of result records as dictionaries
"""
def _execute(tx):
result = tx.run(query, parameters or {})
cleaned_params = self._clean_parameters(parameters)
result = tx.run(query, cleaned_params)
return result.data()
with self._driver.session() as session:
@@ -121,7 +130,8 @@ class Neo4jConnection:
"""
def _execute_batch(tx):
for query, params in queries:
tx.run(query, params or {})
cleaned_params = self._clean_parameters(params)
tx.run(query, cleaned_params)
with self._driver.session() as session:
session.execute_write(_execute_batch)
flowsint-core/src/flowsint_core/core/graph_repository.py
+512 -294
View File
File diff suppressed because it is too large Load Diff
flowsint-core/src/flowsint_core/core/graph_serializer.py
+24 -16
View File
@@ -26,7 +26,7 @@ class GraphSerializer:
- Nested objects (flattened with prefixed keys)
- Lists (converted to primitive types)
- Dictionaries (flattened with prefixed keys)
- None values (converted to empty strings)
- None values (skipped)
Args:
properties: Dictionary of properties to serialize
@@ -37,6 +37,8 @@ class GraphSerializer:
serialized = {}
for key, value in properties.items():
if key is None:
continue
if value is None:
serialized[key] = ""
elif GraphSerializer._is_pydantic_model(value):
@@ -84,21 +86,24 @@ class GraphSerializer:
# Try Pydantic v2 first, then v1
if hasattr(model, "model_dump"):
data = model.model_dump()
data = model.model_dump(mode="json")
elif hasattr(model, "dict"):
data = model.dict()
else:
# Fallback to __dict__
data = model.__dict__
data = {k: v for k, v in model.__dict__.items() if k is not None}
for nested_key, nested_value in data.items():
if nested_value is not None:
new_key = f"{prefix}_{nested_key}"
if isinstance(nested_value, (str, int, float, bool)):
flattened[new_key] = nested_value
else:
# Recursively handle nested complex types
flattened[new_key] = str(nested_value)
if nested_key is None:
continue
new_key = f"{prefix}_{nested_key}"
if nested_value is None:
flattened[new_key] = ""
elif isinstance(nested_value, (str, int, float, bool)):
flattened[new_key] = nested_value
else:
# Recursively handle nested complex types
flattened[new_key] = str(nested_value)
return flattened
@@ -117,12 +122,15 @@ class GraphSerializer:
flattened = {}
for dict_key, dict_value in data.items():
if dict_value is not None:
new_key = f"{prefix}_{dict_key}"
if isinstance(dict_value, (str, int, float, bool)):
flattened[new_key] = dict_value
else:
flattened[new_key] = str(dict_value)
if dict_key is None:
continue
new_key = f"{prefix}_{dict_key}"
if dict_value is None:
flattened[new_key] = ""
elif isinstance(dict_value, (str, int, float, bool)):
flattened[new_key] = dict_value
else:
flattened[new_key] = str(dict_value)
return flattened
flowsint-core/src/flowsint_core/core/graph_service.py
+32 -59
View File
@@ -5,8 +5,9 @@ This module provides a service layer for graph operations,
integrating repository and logging functionality.
"""
from typing import Dict, Any, Optional, Protocol
from typing import Dict, Any, Optional, Protocol, Union
from uuid import UUID
from pydantic import BaseModel
from .graph_repository import GraphRepository
from .graph_db import Neo4jConnection
@@ -33,7 +34,7 @@ class GraphService:
sketch_id: str,
neo4j_connection: Optional[Neo4jConnection] = None,
logger: Optional[LoggerProtocol] = None,
enable_batching: bool = True
enable_batching: bool = True,
):
"""
Initialize the graph service.
@@ -59,94 +60,66 @@ class GraphService:
"""Get the underlying repository."""
return self._repository
def create_node(
self,
node_type: str,
key_prop: str,
key_value: str,
**properties: Any
) -> None:
def create_node(self, node_obj: BaseModel, **properties: Any) -> None:
"""
Create or update a node in the graph.
Automatically adds the following properties:
- type: Lowercase version of node_type
- sketch_id: Current sketch ID
- label: Defaults to key_value if not provided
- created_at: ISO 8601 UTC timestamp (only on creation via ON CREATE SET)
Supports one signatures:
- Pydantic object: create_node(obj, **overrides)
Args:
node_type: Node label (e.g., "domain", "ip")
key_prop: Property name used as unique identifier
key_value: Value of the key property
**properties: Additional node properties
from_obj: a Pydantic object
**properties: Additional node properties or overrides
"""
if self._enable_batching:
self._repository.add_to_batch(
"node",
node_type=node_type,
key_prop=key_prop,
key_value=key_value,
node_obj=node_obj,
sketch_id=self._sketch_id,
**properties
**properties,
)
else:
self._repository.create_node(
node_type=node_type,
key_prop=key_prop,
key_value=key_value,
node_obj=node_obj,
sketch_id=self._sketch_id,
**properties
**properties,
)
def create_relationship(
self,
from_type: str,
from_key: str,
from_value: str,
to_type: str,
to_key: str,
to_value: str,
rel_type: str,
**properties: Any
from_obj: BaseModel,
to_obj: BaseModel,
rel_label: Optional[str] = None,
**properties: Any,
) -> None:
"""
Create a relationship between two nodes.
Supports 1 signature:
- Pydantic objects: create_relationship(obj1, obj2, "REL_TYPE")
Args:
from_type: Source node label
from_key: Source node key property
from_value: Source node key value
to_type: Target node label
to_key: Target node key property
to_value: Target node key value
rel_type: Relationship type
from_obj: Either a Pydantic object (source)
to_obj: Either a Pydantic object (target)
rel_label: Relationship label (ex: "IS_CONNECTED_TO")
**properties: Additional relationship properties
"""
if self._enable_batching:
self._repository.add_to_batch(
"relationship",
from_type=from_type,
from_key=from_key,
from_value=from_value,
to_type=to_type,
to_key=to_key,
to_value=to_value,
rel_type=rel_type,
from_obj=from_obj,
to_obj=to_obj,
rel_label=rel_label,
sketch_id=self._sketch_id,
**properties
**properties,
)
else:
self._repository.create_relationship(
from_type=from_type,
from_key=from_key,
from_value=from_value,
to_type=to_type,
to_key=to_key,
to_value=to_value,
rel_type=rel_type,
from_obj=from_obj,
to_obj=to_obj,
rel_label=rel_label,
sketch_id=self._sketch_id,
**properties
**properties,
)
def log_graph_message(self, message: str) -> None:
@@ -199,7 +172,7 @@ class GraphService:
def create_graph_service(
sketch_id: str,
neo4j_connection: Optional[Neo4jConnection] = None,
enable_batching: bool = True
enable_batching: bool = True,
) -> GraphService:
"""
Factory function to create a GraphService instance.
@@ -219,5 +192,5 @@ def create_graph_service(
sketch_id=sketch_id,
neo4j_connection=neo4j_connection,
logger=Logger,
enable_batching=enable_batching
enable_batching=enable_batching,
)
flowsint-core/tests/enrichers/base.py → flowsint-core/tests/core/base.py
View File
flowsint-core/tests/core/test_enricher_base_simplified_api.py
+52
View File
@@ -0,0 +1,52 @@
"""Test simplified API for create_node and create_relationship."""
import pytest
from flowsint_core.core.enricher_base import Enricher
from flowsint_types.domain import Domain
from flowsint_types.email import Email
from flowsint_types.individual import Individual
from typing import List
class MockEnricher(Enricher):
"""Simple enricher for testing."""
InputType = Domain
OutputType = Domain
@classmethod
def name(cls) -> str:
return "test_enricher"
@classmethod
def category(cls) -> str:
return "Test"
@classmethod
def key(cls) -> str:
return "domain"
async def scan(self, data: List[InputType]) -> List[OutputType]:
return data
def test_create_relationship_with_pydantic_objects():
"""Test that create_relationship works with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Create objects
individual = Individual(first_name="John", last_name="Doe", full_name="John Doe")
domain = Domain(domain="example.com")
# This should not raise an error
enricher.create_relationship(individual, domain, "HAS_DOMAIN")
def test_create_node_with_property_override():
"""Test that property overrides work with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
domain = Domain(domain="example.com")
# Should be able to override properties
enricher.create_node(domain, type="subdomain")
flowsint-core/tests/core/test_graph_repository.py
+260
View File
@@ -0,0 +1,260 @@
"""Test GraphRepository batch operations."""
import pytest
from unittest.mock import Mock, MagicMock
from flowsint_core.core.graph_repository import GraphRepository
from flowsint_types.domain import Domain
from flowsint_types.ip import Ip
from flowsint_types.email import Email
@pytest.fixture
def mock_connection():
"""Create a mock Neo4j connection."""
conn = Mock()
conn.execute_batch = Mock(return_value=[])
conn.query = Mock(return_value=[])
return conn
@pytest.fixture
def graph_repo(mock_connection):
"""Create a GraphRepository with mocked connection."""
return GraphRepository(neo4j_connection=mock_connection)
def test_batch_create_nodes_empty_list(graph_repo):
"""Test batch_create_nodes with empty list."""
result = graph_repo.batch_create_nodes(nodes=[], sketch_id="test-sketch")
assert result["nodes_created"] == 0
assert result["node_ids"] == []
assert result["errors"] == []
def test_batch_create_nodes_single_node(graph_repo, mock_connection):
"""Test batch_create_nodes with a single node."""
domain = Domain(domain="example.com")
# Mock the batch execution to return a node ID
mock_connection.execute_batch.return_value = [
[{"id": "element-id-123"}]
]
result = graph_repo.batch_create_nodes(
nodes=[domain],
sketch_id="test-sketch"
)
assert result["nodes_created"] == 1
assert len(result["node_ids"]) == 1
assert result["node_ids"][0] == "element-id-123"
assert result["errors"] == []
# Verify execute_batch was called once
assert mock_connection.execute_batch.call_count == 1
def test_batch_create_nodes_multiple_nodes(graph_repo, mock_connection):
"""Test batch_create_nodes with multiple nodes of different types."""
domain = Domain(domain="example.com")
ip = Ip(address="192.168.1.1")
email = Email(email="test@example.com")
# Mock the batch execution to return multiple node IDs
mock_connection.execute_batch.return_value = [
[{"id": "element-id-1"}],
[{"id": "element-id-2"}],
[{"id": "element-id-3"}]
]
result = graph_repo.batch_create_nodes(
nodes=[domain, ip, email],
sketch_id="test-sketch"
)
assert result["nodes_created"] == 3
assert len(result["node_ids"]) == 3
assert result["errors"] == []
# Verify execute_batch was called once with 3 operations
assert mock_connection.execute_batch.call_count == 1
batch_operations = mock_connection.execute_batch.call_args[0][0]
assert len(batch_operations) == 3
def test_batch_create_nodes_with_validation_errors(graph_repo, mock_connection):
"""Test batch_create_nodes when some nodes have validation errors."""
valid_domain = Domain(domain="example.com")
# Mock execute_batch to succeed for valid nodes
mock_connection.execute_batch.return_value = [
[{"id": "element-id-1"}]
]
result = graph_repo.batch_create_nodes(
nodes=[valid_domain],
sketch_id="test-sketch"
)
assert result["nodes_created"] == 1
assert len(result["errors"]) == 0
def test_batch_create_nodes_batch_execution_failure(graph_repo, mock_connection):
"""Test batch_create_nodes when the batch execution fails."""
domain = Domain(domain="example.com")
# Mock execute_batch to raise an exception
mock_connection.execute_batch.side_effect = Exception("Database connection error")
result = graph_repo.batch_create_nodes(
nodes=[domain],
sketch_id="test-sketch"
)
assert result["nodes_created"] == 0
assert result["node_ids"] == []
assert len(result["errors"]) > 0
assert "Database connection error" in result["errors"][0]
def test_batch_create_nodes_no_connection():
"""Test batch_create_nodes when there's no database connection."""
# Create a repository with no connection by manually setting it to None
repo = GraphRepository(neo4j_connection=Mock())
repo._connection = None # Force no connection
domain = Domain(domain="example.com")
result = repo.batch_create_nodes(
nodes=[domain],
sketch_id="test-sketch"
)
assert result["nodes_created"] == 0
assert result["node_ids"] == []
assert len(result["errors"]) == 1
assert "No database connection" in result["errors"][0]
def test_batch_create_nodes_with_label_fallback(graph_repo, mock_connection):
"""Test batch_create_nodes with nodes that need label fallback."""
from pydantic import BaseModel
class CustomNode(BaseModel):
"""Node type without a clear primary field."""
name: str = ""
label: str = "Custom Node"
node = CustomNode(name="", label="Fallback Label")
# Mock the batch execution
mock_connection.execute_batch.return_value = [
[{"id": "element-id-fallback"}]
]
result = graph_repo.batch_create_nodes(
nodes=[node],
sketch_id="test-sketch"
)
# Should succeed with fallback to label
assert result["nodes_created"] == 1
assert len(result["node_ids"]) == 1
def test_batch_create_nodes_large_batch(graph_repo, mock_connection):
"""Test batch_create_nodes with a large number of nodes."""
# Create 1000 domains
domains = [Domain(domain=f"example{i}.com") for i in range(1000)]
# Mock the batch execution to return 1000 node IDs
mock_connection.execute_batch.return_value = [
[{"id": f"element-id-{i}"}] for i in range(1000)
]
result = graph_repo.batch_create_nodes(
nodes=domains,
sketch_id="test-sketch"
)
assert result["nodes_created"] == 1000
assert len(result["node_ids"]) == 1000
assert result["errors"] == []
# Verify execute_batch was called once (single transaction)
assert mock_connection.execute_batch.call_count == 1
batch_operations = mock_connection.execute_batch.call_args[0][0]
assert len(batch_operations) == 1000
# Tests for update_node
def test_update_node_success(graph_repo, mock_connection):
"""Test update_node with a valid Pydantic object."""
domain = Domain(domain="updated-example.com", label="Updated Domain")
# Mock the query to return the element ID
mock_connection.query.return_value = [{"id": "element-id-123"}]
result = graph_repo.update_node(
element_id="element-id-123",
node_obj=domain,
sketch_id="test-sketch"
)
assert result == "element-id-123"
assert mock_connection.query.call_count == 1
def test_update_node_not_found(graph_repo, mock_connection):
"""Test update_node when node doesn't exist."""
domain = Domain(domain="example.com")
# Mock the query to return empty result
mock_connection.query.return_value = []
result = graph_repo.update_node(
element_id="non-existent-id",
node_obj=domain,
sketch_id="test-sketch"
)
assert result is None
def test_update_node_no_connection():
"""Test update_node when there's no database connection."""
repo = GraphRepository(neo4j_connection=Mock())
repo._connection = None
domain = Domain(domain="example.com")
result = repo.update_node(
element_id="element-id-123",
node_obj=domain,
sketch_id="test-sketch"
)
assert result is None
def test_update_node_different_types(graph_repo, mock_connection):
"""Test update_node with different Pydantic types."""
ip = Ip(address="10.0.0.1", label="Updated IP")
# Mock the query
mock_connection.query.return_value = [{"id": "element-id-456"}]
result = graph_repo.update_node(
element_id="element-id-456",
node_obj=ip,
sketch_id="test-sketch"
)
assert result == "element-id-456"
# Verify the query contains the correct type
query_call = mock_connection.query.call_args
params = query_call[0][1]
assert params["type"] == "ip"
flowsint-core/tests/test_logger_singleton.py → flowsint-core/tests/core/test_logger_singleton.py
+5 -3
View File
@@ -35,9 +35,11 @@ def mock_db_session():
@pytest.fixture
def mock_get_db(mock_db_session):
"""Mock the get_db generator."""
with patch('flowsint_core.core.logger.get_db') as mock:
mock.return_value = iter([mock_db_session])
yield mock
def mock_generator():
yield mock_db_session
with patch('flowsint_core.core.logger.get_db', mock_generator):
yield mock_generator
@pytest.fixture
flowsint-core/tests/enrichers/test_enricher_base_simplified_api.py
-110
View File
@@ -1,110 +0,0 @@
"""Test simplified API for create_node and create_relationship."""
import pytest
from flowsint_core.core.enricher_base import Enricher
from flowsint_types.domain import Domain
from flowsint_types.email import Email
from flowsint_types.individual import Individual
from typing import List
class MockEnricher(Enricher):
"""Simple enricher for testing."""
InputType = Domain
OutputType = Domain
@classmethod
def name(cls) -> str:
return "test_enricher"
@classmethod
def category(cls) -> str:
return "Test"
@classmethod
def key(cls) -> str:
return "domain"
async def scan(self, data: List[InputType]) -> List[OutputType]:
return data
def test_create_node_with_pydantic_object():
"""Test that create_node works with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Create a domain object
domain = Domain(domain="example.com")
# This should not raise an error
enricher.create_node(domain)
# Verify the helper method works
primary_field = enricher._get_primary_field(domain)
assert primary_field == "domain"
def test_create_relationship_with_pydantic_objects():
"""Test that create_relationship works with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Create objects
individual = Individual(first_name="John", last_name="Doe", full_name="John Doe")
domain = Domain(domain="example.com")
# This should not raise an error
enricher.create_relationship(individual, domain, "HAS_DOMAIN")
def test_create_node_legacy_signature():
"""Test that legacy create_node signature still works."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
domain = Domain(domain="example.com")
# Legacy signature should still work
enricher.create_node("domain", "domain", "example.com", **domain.__dict__)
def test_create_relationship_legacy_signature():
"""Test that legacy create_relationship signature still works."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Legacy signature should still work
enricher.create_relationship(
"individual",
"full_name",
"John Doe",
"domain",
"domain",
"example.com",
"HAS_DOMAIN",
)
def test_get_primary_field():
"""Test the _get_primary_field helper method."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Test with Domain (has primary field marked)
domain = Domain(domain="example.com")
assert enricher._get_primary_field(domain) == "domain"
# Test with Email (has primary field marked)
email = Email(email="test@example.com")
assert enricher._get_primary_field(email) == "email"
# Test with Individual (no primary field marked, falls back to first required field)
individual = Individual(first_name="John", last_name="Doe", full_name="John Doe")
assert enricher._get_primary_field(individual) == "full_name"
def test_create_node_with_property_override():
"""Test that property overrides work with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
domain = Domain(domain="example.com")
# Should be able to override properties
enricher.create_node(domain, type="subdomain")
flowsint-core/tests/test_enricher_base_simplified_api.py
-117
View File
@@ -1,117 +0,0 @@
"""Test simplified API for create_node and create_relationship."""
import pytest
from flowsint_core.core.enricher_base import Enricher
from flowsint_types.domain import Domain
from flowsint_types.email import Email
from flowsint_types.individual import Individual
from typing import List
class MockEnricher(Enricher):
"""Simple enricher for testing."""
InputType = Domain
OutputType = Domain
@classmethod
def name(cls) -> str:
return "test_enricher"
@classmethod
def category(cls) -> str:
return "Test"
@classmethod
def key(cls) -> str:
return "domain"
async def scan(self, data: List[InputType]) -> List[OutputType]:
return data
def test_create_node_with_pydantic_object():
"""Test that create_node works with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Create a domain object
domain = Domain(domain="example.com")
# This should not raise an error
enricher.create_node(domain)
# Verify the helper method works
primary_field = enricher._get_primary_field(domain)
assert primary_field == "domain"
def test_create_relationship_with_pydantic_objects():
"""Test that create_relationship works with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Create objects
individual = Individual(
first_name="John",
last_name="Doe",
full_name="John Doe"
)
domain = Domain(domain="example.com")
# This should not raise an error
enricher.create_relationship(individual, domain, "HAS_DOMAIN")
def test_create_node_legacy_signature():
"""Test that legacy create_node signature still works."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
domain = Domain(domain="example.com")
# Legacy signature should still work
enricher.create_node("domain", "domain", "example.com", **domain.__dict__)
def test_create_relationship_legacy_signature():
"""Test that legacy create_relationship signature still works."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Legacy signature should still work
enricher.create_relationship(
"individual",
"full_name",
"John Doe",
"domain",
"domain",
"example.com",
"HAS_DOMAIN"
)
def test_get_primary_field():
"""Test the _get_primary_field helper method."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
# Test with Domain (has primary field marked)
domain = Domain(domain="example.com")
assert enricher._get_primary_field(domain) == "domain"
# Test with Email (has primary field marked)
email = Email(email="test@example.com")
assert enricher._get_primary_field(email) == "email"
# Test with Individual (no primary field marked, falls back to first required field)
individual = Individual(
first_name="John",
last_name="Doe",
full_name="John Doe"
)
assert enricher._get_primary_field(individual) == "full_name"
def test_create_node_with_property_override():
"""Test that property overrides work with Pydantic objects."""
enricher = MockEnricher(sketch_id="test", scan_id="test")
domain = Domain(domain="example.com")
# Should be able to override properties
enricher.create_node(domain, type="subdomain")