b4w data factories and tests

src/main.py

@@ -3,9 +3,9 @@
 import argparse
 from pathlib import Path
 
-from src.reader import reader_b4w, reader_raw
-from src.helper_db import upsert_b4w
 from src.db import get_session
+from src.helper_db import upsert_b4w
+from src.reader import reader_b4w
 
 
 def ingest_b4w(path: Path):

src/models.py

@@ -4,7 +4,7 @@ from datetime import date
 from decimal import Decimal
 from enum import Enum
 
-from pydantic import BaseModel, Field, ValidationError, field_validator
+from pydantic import BaseModel, Field, field_validator
 
 IBAN_PATTERN = r"^[A-Z]{2}\d{2}[A-Z0-9]{11,30}$"
 
@@ -57,12 +57,16 @@ class B4wExport(BaseModel):
 
     @field_validator("category", mode="before")
     @classmethod
-    def category_str_is_parsable(cls, category_str: str) -> Category | None:
-        category = Category.from_str(category_str)
-        assert category
-        category_str_2 = category.to_str()
-        if not category_str == category_str_2:
-            raise ValidationError(f"Category string cannot be parsed: {category_str}")
+    def category_str_is_parsable(
+        cls, category_in: str | Category | None
+    ) -> Category | None:
+        if category_in is None or category_in == "":
+            return None
+        if isinstance(category_in, Category):
+            return category_in
+        category = Category.from_str(category_in)
+        if category is None or category.to_str() != category_in:
+            raise ValueError(f"Category string cannot be parsed: {category_in!r}")
         return category
 
 

tests/factories.py

@@ -1,13 +1,10 @@
-#!/bin/python
+from collections.abc import Callable
 
 from polyfactory.factories.pydantic_factory import ModelFactory
 
 from src.models import B4wExport
 
-CATEGORIES = [
-        "single_category",
-        "root_category:parent_category:leaf_category",
-        ]
+DEFAULT_CATEGORY = "root_category:parent_category:leaf_category"
 
 class B4wExportFactory(ModelFactory[B4wExport]):
-    __model__ = B4wExport
+    ...

tests/test_model_b4wexport.py

@@ -0,0 +1,12 @@
+from src.models import B4wExport, Category
+from tests.factories import B4wExportFactory
+
+def test_b4wexport_simple():
+    data = B4wExportFactory.build()
+    assert isinstance(data, B4wExport)
+
+
+def test_b4export_category():
+    data = B4wExportFactory.build(Category="fun:stuff")
+    assert isinstance(data.category, Category)
+    assert data.category.to_str() == "fun:stuff"