Fix for list_workbaskets to accept multiple workbasket ids.

exporter/management/commands/dump_transactions.py

@@ -1,22 +1,16 @@
-import os
+import ast
+import itertools
 import sys
 
 from django.conf import settings
 from django.core.management import BaseCommand
 from django.db.transaction import atomic
-from lxml import etree
 
-from common.serializers import validate_envelope
-from exporter.serializers import MultiFileEnvelopeTransactionSerializer
-from exporter.util import dit_file_generator
-from exporter.util import item_timer
+from exporter.management.commands.util import dump_transactions
 from taric.models import Envelope
 from workbaskets.models import WorkBasket
 from workbaskets.validators import WorkflowStatus
 
-# VARIATION_SELECTOR enables emoji presentation
-WARNING_SIGN_EMOJI = "\N{WARNING SIGN}\N{VARIATION SELECTOR-16}"
-
 
 class Command(BaseCommand):
     """
@@ -52,7 +46,7 @@ def add_arguments(self, parser):
                 "with a comma-separated list of workbasket ids."
             ),
             nargs="*",
-            type=int,
+            type=ast.literal_eval,
             default=None,
             action="store",
         )
@@ -76,7 +70,7 @@ def add_arguments(self, parser):
     def handle(self, *args, **options):
         workbasket_ids = options.get("workbasket_ids")
         if workbasket_ids:
-            query = dict(id__in=workbasket_ids)
+            query = dict(id__in=itertools.chain.from_iterable(workbasket_ids))
         else:
             query = dict(status=WorkflowStatus.APPROVED)
 
@@ -106,34 +100,12 @@ def handle(self, *args, **options):
 
         directory = options.get("directory", ".")
 
-        output_file_constructor = dit_file_generator(directory, envelope_id)
-        serializer = MultiFileEnvelopeTransactionSerializer(
-            output_file_constructor,
-            envelope_id=envelope_id,
-            max_envelope_size=max_envelope_size,
+        success = dump_transactions(
+            transactions,
+            envelope_id,
+            directory,
+            max_envelope_size,
+            self.stdout,
         )
-        errors = False
-        for time_to_render, rendered_envelope in item_timer(
-            serializer.split_render_transactions(transactions),
-        ):
-            envelope_file = rendered_envelope.output
-            if not rendered_envelope.transactions:
-                self.stdout.write(
-                    f"{envelope_file.name} {WARNING_SIGN_EMOJI}  is empty !",
-                )
-                errors = True
-            else:
-                envelope_file.seek(0, os.SEEK_SET)
-                try:
-                    validate_envelope(envelope_file)
-                except etree.DocumentInvalid:
-                    self.stdout.write(
-                        f"{envelope_file.name} {WARNING_SIGN_EMOJI}️ Envelope invalid:",
-                    )
-                else:
-                    total_transactions = len(rendered_envelope.transactions)
-                    self.stdout.write(
-                        f"{envelope_file.name} \N{WHITE HEAVY CHECK MARK}  XML valid.  {total_transactions} transactions, serialized in {time_to_render:.2f} seconds using {envelope_file.tell()} bytes.",
-                    )
-        if errors:
+        if success:
             sys.exit(1)

exporter/management/commands/util.py

@@ -0,0 +1,70 @@
+import os
+import sys
+from typing import Optional
+from typing import Sequence
+
+from lxml import etree
+
+from common.models import Transaction
+from common.serializers import validate_envelope
+from exporter.serializers import MultiFileEnvelopeTransactionSerializer
+from exporter.util import dit_file_generator
+from exporter.util import item_timer
+
+# VARIATION_SELECTOR enables emoji presentation
+WARNING_SIGN_EMOJI = "\N{WARNING SIGN}\N{VARIATION SELECTOR-16}"
+
+
+def dump_transactions(
+    transactions: Sequence[Transaction],
+    envelope_id: int,
+    directory: str,
+    max_envelope_size: Optional[int],
+    output_stream=None,
+):
+    """
+    Dump transactions to envelope files in specified directory.
+
+    :param transactions:  Transactions to be pack into envelopes
+    :param envelope_id:  First envelope id to use
+    :param directory:  Directory to write envelope files to.
+    :param max_envelope_size:  Maximum envelope size in bytes, or None to disable splitting.
+    :param output_stream:  Stream to write output status messages to, defaults to stdout.
+
+    See `EnvelopeSerializer` for more information on splitting by size.
+    """
+
+    if output_stream is None:
+        output_stream = sys.stdout
+
+    output_file_constructor = dit_file_generator(directory, envelope_id)
+    serializer = MultiFileEnvelopeTransactionSerializer(
+        output_file_constructor,
+        envelope_id=envelope_id,
+        max_envelope_size=max_envelope_size,
+    )
+    errors = False
+    for time_to_render, rendered_envelope in item_timer(
+        serializer.split_render_transactions(transactions),
+    ):
+        envelope_file = rendered_envelope.output
+        if not rendered_envelope.transactions:
+            output_stream.write(
+                f"{envelope_file.name} {WARNING_SIGN_EMOJI}  is empty !",
+            )
+            errors = True
+        else:
+            envelope_file.seek(0, os.SEEK_SET)
+            try:
+                validate_envelope(envelope_file)
+                assert 0
+            except etree.DocumentInvalid:
+                output_stream.write(
+                    f"{envelope_file.name} {WARNING_SIGN_EMOJI}️ Envelope invalid:",
+                )
+            else:
+                total_transactions = len(rendered_envelope.transactions)
+                output_stream.write(
+                    f"{envelope_file.name} \N{WHITE HEAVY CHECK MARK} XML valid. {total_transactions} transactions, serialized in {time_to_render:.2f} seconds using {envelope_file.tell()} bytes.",
+                )
+    return not errors

exporter/tests/test_commands_util.py

@@ -0,0 +1,65 @@
+from unittest.mock import patch
+
+import pytest
+from lxml import etree
+
+from common.tests import factories
+from exporter.management.commands import util
+
+pytestmark = pytest.mark.django_db
+
+
+def test_dump_transactions_success(tmp_path, capsys, approved_transaction):
+    factories.GoodsNomenclatureFactory.create(transaction=approved_transaction)
+    transactions = approved_transaction.workbasket.transactions.all()
+    envelope_id = 123
+    directory = tmp_path / "test_directory"
+    directory.mkdir()
+    util.dump_transactions(transactions, envelope_id, directory, None)
+    output = capsys.readouterr().out
+    path = directory.__str__()
+    total_transactions = approved_transaction.workbasket.transactions.count()
+
+    assert (
+        f"{path}/DIT000123.xml \N{WHITE HEAVY CHECK MARK} XML valid. {total_transactions} transactions, serialized"
+        in output
+    )
+
+
+def test_dump_transactions_empty(tmp_path, capsys, approved_transaction):
+    transactions = approved_transaction.workbasket.transactions.all()
+    envelope_id = 123
+    directory = tmp_path / "test_directory"
+    directory.mkdir()
+    util.dump_transactions(transactions, envelope_id, directory, None)
+    output = capsys.readouterr().out
+    path = directory.__str__()
+
+    assert (
+        f"{path}/DIT000123.xml \N{WARNING SIGN}\N{VARIATION SELECTOR-16}  is empty !"
+        == output
+    )
+
+
+@patch("common.serializers.validate_envelope")
+def test_dump_transactions_invalid(
+    validate_envelope,
+    tmp_path,
+    capsys,
+    approved_transaction,
+):
+    validate_envelope.side_effect = etree.DocumentInvalid("invalid")
+    assert 0
+    factories.GoodsNomenclatureFactory.create(transaction=approved_transaction)
+    transactions = approved_transaction.workbasket.transactions.all()
+    envelope_id = 123
+    directory = tmp_path / "test_directory"
+    directory.mkdir()
+    util.dump_transactions(transactions, envelope_id, directory, None)
+    output = capsys.readouterr().out
+    path = directory.__str__()
+
+    assert (
+        f"{path}/DIT000123.xml \N{WARNING SIGN}\N{VARIATION SELECTOR-16}  Envelope invalid:"
+        == output
+    )

pii-ner-exclude.txt

@@ -1216,3 +1216,5 @@ f"{first_line_of(workbasket.reason
 WorkBasketOutputFormat Enum
 Enter the Regulation ID
 Enter the Quota Order Number
+f"{path}/DIT000123.xml \N{WARNING SIGN}\N{VARIATION SELECTOR-16
+f"{path}/DIT000123.xml \N{WHITE HEAVY

workbaskets/management/commands/list_workbaskets.py

@@ -56,7 +56,8 @@ def add_arguments(self, parser: CommandParser) -> None:
 
         parser.add_argument(
             "workbasket_ids",
-            help=("Comma-separated list of workbasket ids to filter to"),
+            nargs="*",
+            help="Comma-separated list of workbasket ids to filter to",
             type=ast.literal_eval,
         )
 

workbaskets/management/util.py

@@ -30,8 +30,14 @@ def _output_workbasket_readable(
             f"{spaces}tracked model count: {workbasket.tracked_models.count()}",
         )
         if show_transaction_info:
+            transactions = workbasket.transactions
+            first_pk = (
+                workbasket.transactions.first().pk if transactions.count() else "-"
+            )
+            last_pk = workbasket.transactions.last().pk if transactions.count() else "-"
+
             self.stdout.write(
-                f"{spaces}transactions: {workbasket.transactions.first().pk} - {workbasket.transactions.last().pk}",
+                f"{spaces}transactions: {first_pk} - {last_pk} [{transactions.count()}]",
             )
 
     def _output_workbasket_compact(self, workbasket, show_transaction_info, **kwargs):
@@ -43,14 +49,19 @@ def _output_workbasket_compact(self, workbasket, show_transaction_info, **kwargs
             ending="" if show_transaction_info else "\n",
         )
         if show_transaction_info:
+            transactions = workbasket.transactions
+            first_pk = (
+                workbasket.transactions.first().pk if transactions.count() else "-"
+            )
+            last_pk = workbasket.transactions.last().pk if transactions.count() else "-"
             self.stdout.write(
-                f", {workbasket.transactions.first().pk} - {workbasket.transactions.last().pk}",
+                f", {first_pk} - {last_pk} [{transactions.count()}]",
             )
 
     def output_workbasket(
         self,
         workbasket,
-        show_transaction_info=False,
+        show_transaction_info,
         output_format=WorkBasketOutputFormat.READABLE,
         **kwargs,
     ):