diff --git a/src/cl_sii/dte/data_models.py b/src/cl_sii/dte/data_models.py index 799711d0..d4cbd1de 100644 --- a/src/cl_sii/dte/data_models.py +++ b/src/cl_sii/dte/data_models.py @@ -216,23 +216,25 @@ def random( cls, emisor_rut: Optional[Rut] = None, tipo_dte: TipoDte | Sequence[TipoDte] = tuple(TipoDte), - folio: Optional[int] = None, + folio: int | tuple[int, int] = ( + constants.DTE_FOLIO_FIELD_MIN_VALUE, + constants.DTE_FOLIO_FIELD_MAX_VALUE, + ), ) -> Self: """ Generate random DTE natural key within valid ranges. :param emisor_rut: RUT of the "emisor" of the DTE. If `None`, a random RUT is generated. :param tipo_dte: The kind of DTE. If a sequence is provided, a random one is chosen. - :param folio: The sequential number of the DTE. If `None`, a random folio is generated. + :param folio: The sequential number of the DTE. If a 2-tuple of integers is provided, + a random one is chosen within the range defined by the tuple. """ if emisor_rut is None: emisor_rut = Rut.random() if isinstance(tipo_dte, Sequence): tipo_dte = random.choice(tipo_dte) - if folio is None: - folio = random.randint( - constants.DTE_FOLIO_FIELD_MIN_VALUE, constants.DTE_FOLIO_FIELD_MAX_VALUE - ) + if isinstance(folio, tuple): + folio = random.randint(*folio) return cls(emisor_rut, tipo_dte, folio) diff --git a/src/tests/test_dte_data_models.py b/src/tests/test_dte_data_models.py index 32d03b30..30ac791c 100644 --- a/src/tests/test_dte_data_models.py +++ b/src/tests/test_dte_data_models.py @@ -138,6 +138,17 @@ def test_random(self) -> None: specific_dte_nk = DteNaturalKey.random(tipo_dte=specific_tipo) self.assertEqual(specific_dte_nk.tipo_dte, specific_tipo) + # Test with sequence of tipo_dte values + tipo_dte_sequence = [TipoDte.FACTURA_ELECTRONICA, TipoDte.NOTA_CREDITO_ELECTRONICA] + seq_dte_nk = DteNaturalKey.random(tipo_dte=tipo_dte_sequence) + self.assertIn(seq_dte_nk.tipo_dte, tipo_dte_sequence) + + # Test with custom folio range + custom_folio_range = (1000, 2000) + range_dte_nk = DteNaturalKey.random(folio=custom_folio_range) + self.assertGreaterEqual(range_dte_nk.folio, custom_folio_range[0]) + self.assertLessEqual(range_dte_nk.folio, custom_folio_range[1]) + # Test with partial custom parameters partial_custom = DteNaturalKey.random(emisor_rut=custom_rut) self.assertEqual(partial_custom.emisor_rut, custom_rut)