bb-integrations-library 3.0.11__py3-none-any.whl

bb_integrations_lib/gravitate/testing/builder.py

@@ -0,0 +1,55 @@
+import time
+
+from loguru import logger
+
+from bb_integrations_lib.gravitate.testing.util import generate_model_validation_tests, \
+    generate_pydantic_models_from_open_api
+
+
+class TestBuilder:
+    def __init__(self,
+                 client: str,
+                 system: str):
+        self.client = client
+        self.system = system
+
+    @property
+    def models_file_path(self):
+        return f"./{self.client}/{self.system}/models.py"
+
+    @property
+    def tests_file_path(self):
+        return f"./{self.client}/{self.system}/tests/test_models.py"
+
+    @property
+    def open_api_url(self):
+        urls = {
+            "sd": f"https:/{self.client}.bb.gravitate.energy/api/openapi.json/internal",
+            "rita": "https://rita.gravitate.energy/api/openapi.json"
+        }
+        return urls[self.system]
+
+    def build_tests(self):
+        try:
+            generate_pydantic_models_from_open_api(
+                open_api_url=self.open_api_url,
+                save_file_to_path=self.models_file_path
+            )
+            time.sleep(1)
+            generate_model_validation_tests(
+                models_file_path=self.models_file_path,
+                tests_file_path=self.tests_file_path
+            )
+        except Exception as e:
+            logger.error(f"Failed to build tests: {e}")
+            raise
+
+    @classmethod
+    def for_client_and_system(cls, client: str, system: str):
+        instance = cls(client=client, system=system)
+        instance.build_tests()
+        return instance
+
+
+if __name__ == "__main__":
+    test_builder = TestBuilder.for_client_and_system("tte", "sd")

bb_integrations_lib/gravitate/testing/openapi.py

@@ -0,0 +1,70 @@
+import json
+
+import httpx
+from functools import lru_cache
+
+
+@lru_cache(maxsize=10)
+def get_openapi(url: str) -> dict:
+    """Download OpenAPI schema from URL and return as dict."""
+    response = httpx.get(url)
+    response.raise_for_status()
+    return response.json()
+
+
+def get_updated_supply_and_dispatch_openapi_json(url: str, schemas_to_include: list[str]):
+    """Update the supply and dispatch JSON file with the latest published schemas."""
+    schema = get_openapi(url)
+    fully_filtered = filter_schemas_by_references(filter_based_on_tag(schema, schemas_to_include))
+    return json.dumps(fully_filtered)
+
+
+def filter_based_on_tag(data: dict, criteria: list[str]) -> dict:
+    """Filter OpenAPI paths based on tags."""
+    filtered_data = data.copy()
+    filtered_paths = {}
+
+    for path, methods in data.get("paths", {}).items():
+        matching_methods = {}
+        for method, operation in methods.items():
+            if isinstance(operation, dict):
+                tags = operation.get("tags", [])
+                if any(tag in tags for tag in criteria):
+                    matching_methods[method] = operation
+        if matching_methods:
+            filtered_paths[path] = matching_methods
+
+    filtered_data["paths"] = filtered_paths
+    return filtered_data
+
+
+def filter_schemas_by_references(data: dict) -> dict:
+    filtered_data = data.copy()
+    referenced_schemas = set()
+    for path, methods in data.get("paths", {}).items():
+        for method, operation in methods.items():
+            if isinstance(operation, dict):
+                _collect_schema_references(operation, referenced_schemas)
+    if "components" in data and "schemas" in data["components"]:
+        filtered_data["components"] = data["components"].copy()
+        filtered_data["components"]["schemas"] = {
+            name: schema for name, schema in data["components"]["schemas"].items()
+            if name in referenced_schemas
+        }
+    return filtered_data
+
+
+def _collect_schema_references(obj, referenced_schemas):
+    if isinstance(obj, dict):
+        for key, value in obj.items():
+            if key == "$ref" and isinstance(value, str):
+                if value.startswith("#/components/schemas/"):
+                    schema_name = value.split("/")[-1]
+                    referenced_schemas.add(schema_name)
+            else:
+                _collect_schema_references(value, referenced_schemas)
+    elif isinstance(obj, list):
+        for item in obj:
+            _collect_schema_references(item, referenced_schemas)
+
+

bb_integrations_lib/gravitate/testing/util.py

@@ -0,0 +1,274 @@
+import importlib.util
+import inspect
+import json
+from pathlib import Path
+from typing import get_origin, Union, get_args, Any, Dict
+import time
+from datamodel_code_generator import generate, DataModelType, InputFileType
+from loguru import logger
+from pydantic import BaseModel
+
+from bb_integrations_lib.gravitate.testing.openapi import get_updated_supply_and_dispatch_openapi_json
+
+
+def ensure_directory_exists(file_path: Path) -> None:
+    """Ensure the parent directory of a file path exists."""
+    file_path.parent.mkdir(parents=True, exist_ok=True)
+    return None
+
+
+def generate_pydantic_models_from_open_api(open_api_url: str,
+                                           save_file_to_path: str,
+                                           schemas_to_include: list[str] = None,
+                                           open_api_json: dict = None) -> None:
+    try:
+        openapi_data = open_api_json or get_updated_supply_and_dispatch_openapi_json(
+            open_api_url,
+            schemas_to_include or ['V1', 'V2'])
+        if isinstance(openapi_data, str):
+            try:
+                openapi_json = json.loads(openapi_data)
+                logger.info("Parsed OpenAPI JSON string to dictionary")
+            except json.JSONDecodeError as e:
+                logger.error(f"Failed to parse OpenAPI JSON string: {e}")
+                raise
+        else:
+            openapi_json = openapi_data
+        if not isinstance(openapi_json, dict):
+            raise ValueError(f"Expected dict or JSON string, got {type(openapi_json)}")
+        save_path = Path(save_file_to_path)
+        temp_openapi_file = save_path.parent / "temp_openapi.json"
+        ensure_directory_exists(save_path)
+        ensure_directory_exists(temp_openapi_file)
+        with open(temp_openapi_file, 'w') as f:
+            json.dump(openapi_json, f, indent=2)
+        generate(
+            input_=temp_openapi_file,
+            output=save_path,
+            output_model_type=DataModelType.PydanticV2BaseModel
+        )
+        if temp_openapi_file.exists():
+            temp_openapi_file.unlink()
+        logger.info(f"Successfully generated Pydantic models at {save_file_to_path}")
+    except Exception as e:
+        logger.error(f"Error generating pydantic models: {e}")
+        temp_openapi_file = Path(save_file_to_path).parent / "temp_openapi.json"
+        if temp_openapi_file.exists():
+            temp_openapi_file.unlink()
+        raise
+    return None
+
+
+def generate_model_validation_tests(models_file_path: str, tests_file_path: str) -> None:
+    """Generate validation tests for all Pydantic models in the models file."""
+    try:
+        if not Path(models_file_path).exists():
+            logger.error(f"Models file does not exist: {models_file_path}")
+            return
+        models = load_pydantic_models_from_file(models_file_path)
+        if not models:
+            logger.warning(f"No Pydantic models found in {models_file_path}")
+            return
+        test_data = {}
+        for model_name, model_class in models.items():
+            try:
+                example_input = generate_example_input(model_class)
+                test_data[model_name] = example_input
+            except Exception as e:
+                logger.warning(f"Could not generate example for {model_name}: {e}")
+                test_data[model_name] = {}
+        create_test_file(tests_file_path, models_file_path, test_data)
+
+        logger.info(f"Generated validation tests for {len(models)} models in {tests_file_path}")
+
+    except Exception as e:
+        logger.error(f"Error generating model validation tests: {e}")
+        raise
+
+
+def load_pydantic_models_from_file(file_path: str) -> Dict[str, BaseModel]:
+    models = {}
+    try:
+        spec = importlib.util.spec_from_file_location("models", file_path)
+        if spec is None or spec.loader is None:
+            logger.error(f"Could not load module spec from {file_path}")
+            return models
+
+        models_module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(models_module)
+        for name in dir(models_module):
+            try:
+                obj = getattr(models_module, name)
+                if (inspect.isclass(obj) and
+                        issubclass(obj, BaseModel) and
+                        obj is not BaseModel):
+                    models[name] = obj
+            except Exception as e:
+                logger.debug(f"Skipping {name}: {e}")
+                continue
+
+    except Exception as e:
+        logger.error(f"Error loading models from {file_path}: {e}")
+
+    return models
+
+
+def generate_example_input(model_class: BaseModel) -> Dict[str, Any]:
+    example = {}
+    try:
+        if hasattr(model_class, 'model_fields'):
+            fields = model_class.model_fields
+            for field_name, field_info in fields.items():
+                try:
+                    example[field_name] = generate_field_example(field_info.annotation, field_name)
+                except Exception as e:
+                    logger.debug(f"Could not generate example for field {field_name}: {e}")
+                    example[field_name] = f"example_{field_name}"
+        elif hasattr(model_class, '__fields__'):
+            fields = model_class.__fields__
+            for field_name, field_info in fields.items():
+                try:
+                    example[field_name] = generate_field_example(field_info.type_, field_name)
+                except Exception as e:
+                    logger.debug(f"Could not generate example for field {field_name}: {e}")
+                    example[field_name] = f"example_{field_name}"
+        else:
+            logger.warning(f"Could not find fields for model {model_class.__name__}")
+    except Exception as e:
+        logger.warning(f"Error generating example input for {model_class.__name__}: {e}")
+
+    return example
+
+
+def generate_field_example(field_type: Any, field_name: str) -> Any:
+    """Generate example value for a field type."""
+    try:
+        origin = get_origin(field_type)
+        args = get_args(field_type)
+        if origin is Union:
+            non_none_types = [arg for arg in args if arg is not type(None)]
+            if non_none_types:
+                field_type = non_none_types[0]
+        if origin is list:
+            if args:
+                item_type = args[0]
+                return [generate_field_example(item_type, f"{field_name}_item")]
+            return ["example_item"]
+        if origin is dict:
+            return {"key": "value"}
+        if field_type == str:
+            return f"example_{field_name}"
+        elif field_type == int:
+            return 42
+        elif field_type == float:
+            return 3.14
+        elif field_type == bool:
+            return True
+        elif field_type == list:
+            return ["example_item"]
+        elif field_type == dict:
+            return {"key": "value"}
+        else:
+            return f"example_{field_name}"
+    except Exception as e:
+        logger.debug(f"Error generating field example for {field_name}: {e}")
+        return f"example_{field_name}"
+
+
+def create_test_file(tests_file_path: str, models_file_path: str, test_data: Dict[str, Dict]) -> None:
+    """Create the test file with validation tests for all models."""
+    tests_path = Path(tests_file_path)
+    ensure_directory_exists(tests_path)
+    models_import_path = Path(models_file_path).stem
+    valid_test_data = {k: v for k, v in test_data.items() if v}
+    if not valid_test_data:
+        logger.warning("No valid test data found, creating minimal test file")
+        test_content = f'''"""
+Auto-generated validation tests for Pydantic models.
+Generated on: {time.strftime('%Y-%m-%d %H:%M:%S')}
+No valid models found for testing.
+"""
+import pytest
+
+def test_placeholder():
+    """Placeholder test - no valid models found."""
+    assert True
+'''
+    else:
+        test_content = f'''"""
+Auto-generated validation tests for Pydantic models.
+Generated on: {time.strftime('%Y-%m-%d %H:%M:%S')}
+"""
+import pytest
+from typing import Dict, Any
+from pydantic import BaseModel, ValidationError
+from {models_import_path} import {", ".join(valid_test_data.keys())}
+
+
+def generic_model_validation_test(model_to_test: BaseModel, example_input: Dict[str, Any]) -> None:
+    """
+    Generic test function that validates a model with example input.
+
+    Args:
+        model_to_test: The Pydantic model class to test
+        example_input: Dictionary with example input data
+    """
+    try:
+        # Test successful validation
+        validated_model = model_to_test.model_validate(example_input)
+        assert validated_model is not None
+
+        # Test that we can convert back to dict
+        model_dict = validated_model.model_dump()
+        assert isinstance(model_dict, dict)
+
+    except ValidationError as e:
+        pytest.fail(f"Validation failed for {{model_to_test.__name__}}: {{e}}")
+    except Exception as e:
+        pytest.fail(f"Unexpected error for {{model_to_test.__name__}}: {{e}}")
+
+
+# Individual test functions for each model
+'''
+
+        # Generate individual test functions for each model
+        for model_name, example_data in valid_test_data.items():
+            test_content += f'''
+def test_{model_name.lower()}_validation():
+    """Test validation for {model_name} model."""
+    example_input = {json.dumps(example_data, indent=4)}
+
+    generic_model_validation_test({model_name}, example_input)
+
+
+def test_{model_name.lower()}_validation_with_invalid_data():
+    """Test validation failure for {model_name} model with invalid data."""
+    invalid_input = {{"invalid_field": "should_fail"}}
+
+    with pytest.raises(ValidationError):
+        {model_name}.model_validate(invalid_input)
+'''
+
+        # Add parametrized test for all models
+        if valid_test_data:
+            test_content += f'''
+
+# Parametrized test for all models
+@pytest.mark.parametrize("model_class,example_input", [
+'''
+
+            for model_name, example_data in valid_test_data.items():
+                test_content += f'    ({model_name}, {json.dumps(example_data)}),\n'
+
+            test_content += '''
+])
+def test_all_models_validation(model_class: BaseModel, example_input: Dict[str, Any]):
+    """Parametrized test for all models."""
+    generic_model_validation_test(model_class, example_input)
+'''
+
+    # Write the test file
+    with open(tests_path, 'w') as f:
+        f.write(test_content)
+
+    logger.info(f"Created test file at {tests_file_path}")

bb_integrations_lib/mappers/prices/model.py

@@ -0,0 +1,106 @@
+from enum import Enum
+from typing import Optional, Dict, Union, List
+
+from mergedeep import Strategy
+from pydantic import BaseModel
+
+from bb_integrations_lib.gravitate.model import System
+from bb_integrations_lib.models.rita.mapping import Map
+
+
+class Action(str, Enum):
+    start = "start"
+    stop = "stop"
+    error = "error"
+
+
+class IntegrationType(str, Enum):
+    """Denotes available integration types for a Mapper"""
+    sql = "sql"
+    """Strictly limited to an MS SQL server database"""
+    rita = "rita"
+    """Strictly limited to rita mappings"""
+
+
+class IntegrationMappingConfig(BaseModel):
+    type: IntegrationType | System = None
+    """The integration type to pull mappings from"""
+    external_id_field: Optional[str] = None
+    """External: to Gravitate. Denotes the id field to be used"""
+    gravitate_id_field: Optional[str] = None
+    """The Gravitate id field"""
+    gravitate_name_field: Optional[str] = None
+    external_name_field: Optional[str] = None
+    query: Optional[str] = None
+    """An optional query string"""
+
+
+class EntityConfig(BaseModel):
+    mapping_enabled: Optional[bool] = True
+    mapping_integration: Optional[IntegrationMappingConfig] = {}
+    external_system_integration: Optional[IntegrationMappingConfig] = {}
+
+class PricePublisher(BaseModel):
+    id: Optional[str] = None
+    name: str
+    price_type: str
+    extend_by_days: Optional[int] = None
+
+class PricingStrategy(str, Enum):
+    use_latest = "Use Latest"
+    """Only includes latest price -> may miss intraday changes"""
+    use_prior_to_latest = "Use Prior to Latest"
+    """Includes both the latest and latest minus one"""
+    use_historic ="Use Historic"
+    """Includes up to 10 historic changes for instrument since previous workday"""
+
+    @property
+    def strategy_includes(self) -> int:
+        if self == PricingStrategy.use_latest:
+            return 1
+        if self == PricingStrategy.use_prior_to_latest:
+            return 2
+        if self == PricingStrategy.use_historic:
+            return 10
+
+class PricingIntegrationConfig(BaseModel):
+    environment: str
+    """The customer environment; i.e. TTE"""
+    price_publishers: List[PricePublisher]
+    """The list of price publisher from which to pull prices"""
+    entity_config: Dict[str, EntityConfig] = {}
+    """A key: EntityConfig pair, describing an entity config. i.e. {'products': EntityConfig}"""
+    price_mapper_ttl: Optional[int] = 3600
+    """The ttl cache release for the mapper"""
+    price_mapper_debug_mode: Optional[bool] = False
+    """Debug mode to enable verbose logging"""
+    source_system: Optional[str] = None
+    """The source system from where prices originate"""
+    source_system_id: Optional[str] = None
+    """The source system id from where prices originate"""
+    use_contract_id: bool = True
+    strategy: Optional[PricingStrategy] = PricingStrategy.use_historic
+
+
+
+class Group(BaseModel):
+    name: Optional[str] = None
+    ids: list[str]
+    length: int
+    extra_data: Optional[Union[Dict, Map]] = None
+
+
+class Groups(BaseModel):
+    product_groups: Optional[Dict[str, Group]] = None
+    location_groups: Optional[Dict[str, Group]] = None
+    supplier_groups: Optional[Dict[str, Group]] = None
+    price_publisher_groups: Optional[Dict[str, Group]] = None
+
+
+class PriceMappings(BaseModel):
+    product_mappings: Union[Dict, List]
+    location_mappings: Union[Dict, List]
+    supplier_mappings: Union[Dict, List]
+    price_publishers: Union[Dict, List]
+
+

bb_integrations_lib/mappers/prices/price_mapper.py

@@ -0,0 +1,127 @@
+from typing import Optional, Union, List, Dict
+from collections.abc import Iterable
+from loguru import logger
+
+from bb_integrations_lib.gravitate.rita_api import RitaBackendAPI
+from bb_integrations_lib.mappers.prices.model import PriceMappings, Group, PricingIntegrationConfig,\
+    IntegrationMappingConfig
+from bb_integrations_lib.mappers.prices.util import PricingIntegrationGetter, PricingIntegrationGetterProtocol
+from bb_integrations_lib.models.rita.mapping import Map
+from bb_integrations_lib.provider.sqlserver.client import SQLServerClient
+from bb_integrations_lib.util.cache.custom_ttl_cache import CustomTTLCache
+from bb_integrations_lib.util.cache.protocol import CustomTTLCacheProtocol
+
+
+class PriceMapper:
+    """Mapper instance to get all mapping data for needed for prices.
+    Limited to:
+        - Product, Location, Supplier and Publisher mappings
+    """
+    def __init__(
+            self,
+            config: PricingIntegrationConfig,
+            ttl: int = 3600,
+            debug_mode: bool = False,
+            source_system: Optional[str] = None
+    ):
+        self.ttl = ttl
+        self.config = config
+        self.products_config: IntegrationMappingConfig = self.config.entity_config.get("products").mapping_integration
+        self.locations_config: IntegrationMappingConfig = self.config.entity_config.get("locations").mapping_integration
+        self.suppliers_config: IntegrationMappingConfig = self.config.entity_config.get("suppliers").mapping_integration
+        self.publishers_config: IntegrationMappingConfig = self.config.entity_config.get("publishers").mapping_integration
+        self.products_sql_query = self.products_config.query
+        self.locations_sql_query = self.locations_config.query
+        self.suppliers_sql_query = self.suppliers_config.query
+        self.cache: CustomTTLCacheProtocol = CustomTTLCache(verbose=debug_mode)
+        self.debug_mode = debug_mode
+        self.source_system = source_system
+
+        if debug_mode:
+            logger.debug("DEBUG MODE ON")
+        self.integration_getter: PricingIntegrationGetterProtocol = PricingIntegrationGetter(self.config)
+
+    @property
+    def ttl_cache(self):
+        return self.cache
+
+    async def _get_cached_data(self, cache_key: str, query_func) -> Iterable:
+        """Generic method to retrieve cached data or fetch it using the provided query function"""
+
+        @self.cache.ttl_cache(seconds=self.ttl, cache_key=cache_key)
+        async def _get_data() -> Iterable:
+            return await query_func() or {}
+        return await _get_data()
+
+
+    async def get_entity_mappings(self, key: str, mapping_type: Optional[str] = None, query: Optional[str] = None) -> Iterable:
+        client = self.integration_getter.get_integration_mapping_client_by_entity(key)
+
+        async def query_func():
+            if isinstance(client, SQLServerClient):
+                return client.get_mappings(query)
+            elif isinstance(client, RitaBackendAPI):
+                return await client.get_mappings(source_system=self.source_system,
+                                           mapping_type=mapping_type)
+            return {}
+        return await self._get_cached_data(key, query_func)
+
+
+    async def get_mappings(
+            self,
+            product_key: Optional[str] = "products",
+            location_key: Optional[str] = "locations",
+            supplier_key: Optional[str] = "suppliers",
+            price_publisher_key: Optional[str] = "publishers"
+    ) -> PriceMappings:
+        return PriceMappings(
+            product_mappings=await self.get_entity_mappings(product_key, mapping_type="product", query=self.products_sql_query),
+            location_mappings=await self.get_entity_mappings(location_key, mapping_type="location", query=self.locations_sql_query),
+            supplier_mappings=await self.get_entity_mappings(supplier_key, mapping_type="supplier", query=self.suppliers_sql_query),
+            price_publishers=await self.get_entity_mappings(price_publisher_key, mapping_type="other")
+        )
+
+    @classmethod
+    def group_rows(
+            cls,
+            rows: Union[List[Map], List[Dict]],
+            external_id_field: str,
+            gravitate_id_field: str,
+            name_field: Optional[str] = None,
+            is_rita: Optional[bool] = False,
+    ) -> Dict[
+        str, Group]:
+        grouped = {}
+        if is_rita:
+            rows: List[Map]
+            for row in rows:
+                row_id = row.source_id
+                children = [r.gravitate_id for r in row.children]
+                if row_id not in grouped:
+                    grouped[row_id] = Group(
+                        name=row.gravitate_name,
+                        ids=children,
+                        length=len(children),
+                        extra_data=row
+                    )
+            return grouped
+        else:
+            rows: List[Dict]
+            for row in rows:
+                row_id = row[external_id_field]
+                row_guid = row[gravitate_id_field]
+                row_name = row.get(name_field) if name_field else None
+                if row_id not in grouped:
+                    grouped[str(row_id)] = Group(
+                        name=row_name,
+                        ids=[row_guid],
+                        length=1,
+                        extra_data=row
+                    )
+                else:
+                    group = grouped[row_id]
+                    group.ids.append(row_guid)
+                    group.length += 1
+        return grouped
+
+

bb_integrations_lib/mappers/prices/protocol.py

@@ -0,0 +1,20 @@
+from typing import Protocol, Optional, Iterable, Dict, Any, runtime_checkable
+
+
+class PriceMapperProtocol(Protocol):
+    async def get_product_mappings(self, product_key: Optional[str]) -> Iterable:
+        """Gets product mappings from integration"""
+
+    async def get_location_mappings(self, location_key: Optional[str]) -> Iterable:
+        """Gets location mappings from integration"""
+
+    async def get_supplier_mappings(self, supplier_key: Optional[str]) -> Iterable:
+        """Gets supplier mappings from integration"""
+
+
+@runtime_checkable
+class ExternalPriceMapperIntegration(Protocol):
+    def get_mappings(self, query: Optional[str] = None, source_system: Optional[str] = None,
+                     mapping_type: Optional[str] = None,
+                     params: Optional[Dict[str, Any]] = None):
+        """Gets mappings from integration"""