From 2136eff78bdfa1c25ad84b0adc5bbfb231d968a1 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Wed, 29 Oct 2025 23:33:21 +0000 Subject: [PATCH 1/2] fix `ModelBuilder` for array columns with choices --- piccolo/testing/model_builder.py | 14 ++++++++++---- tests/testing/test_model_builder.py | 15 +++++++++++++++ 2 files changed, 25 insertions(+), 4 deletions(-) diff --git a/piccolo/testing/model_builder.py b/piccolo/testing/model_builder.py index b493f9afa..7e25d7e10 100644 --- a/piccolo/testing/model_builder.py +++ b/piccolo/testing/model_builder.py @@ -173,10 +173,16 @@ def _randomize_attribute(cls, column: Column) -> Any: random_value = RandomBuilder.next_datetime(tz_aware=tz_aware) elif column.value_type == list: length = RandomBuilder.next_int(maximum=10) - base_type = cast(Array, column).base_column.value_type - random_value = [ - cls.__DEFAULT_MAPPER[base_type]() for _ in range(length) - ] + if column._meta.choices: + random_value = [ + RandomBuilder.next_enum(column._meta.choices) + for _ in range(length) + ] + else: + base_type = cast(Array, column).base_column.value_type + random_value = [ + cls.__DEFAULT_MAPPER[base_type]() for _ in range(length) + ] elif column._meta.choices: random_value = RandomBuilder.next_enum(column._meta.choices) else: diff --git a/tests/testing/test_model_builder.py b/tests/testing/test_model_builder.py index dba0dc791..2b708c74d 100644 --- a/tests/testing/test_model_builder.py +++ b/tests/testing/test_model_builder.py @@ -1,4 +1,5 @@ import asyncio +import enum import json import unittest @@ -30,9 +31,14 @@ class TableWithArrayField(Table): + class Choices(enum.Enum): + a = "a" + b = "b" + strings = Array(Varchar(30)) integers = Array(Integer()) floats = Array(Real()) + choices = Array(Varchar(), choices=Choices) class TableWithDecimal(Table): @@ -104,6 +110,15 @@ def test_choices(self): ["s", "l", "m"], ) + def test_array_choices(self): + """ + Make sure that ``ModelBuilder`` generates arrays where each array + element is a valid choice. + """ + instance = ModelBuilder.build_sync(TableWithArrayField) + for value in instance.choices: + assert value in TableWithArrayField.Choices + def test_datetime(self): """ Make sure that ``ModelBuilder`` generates timezone aware datetime From b5f98e942bef6a1389f42e6d46300add57256c62 Mon Sep 17 00:00:00 2001 From: Daniel Townsend Date: Wed, 29 Oct 2025 23:41:13 +0000 Subject: [PATCH 2/2] make test work with older Python versions --- tests/testing/test_model_builder.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/testing/test_model_builder.py b/tests/testing/test_model_builder.py index 2b708c74d..b1d07376a 100644 --- a/tests/testing/test_model_builder.py +++ b/tests/testing/test_model_builder.py @@ -117,7 +117,8 @@ def test_array_choices(self): """ instance = ModelBuilder.build_sync(TableWithArrayField) for value in instance.choices: - assert value in TableWithArrayField.Choices + # Will raise an exception if the enum value isn't found: + TableWithArrayField.Choices[value] def test_datetime(self): """