import unittest.mock

import generator.arrange.form
import generator.arrange.generate
import generator.parts

def section_document(document, part_type, *, rhythms, plan=None):
  """Exercise the real part pipeline with a single section instead of a full song."""
  prepared = generator.parts.prepare_document(document)
  role = {'acoustic': 'accompaniment', 'electric': 'melody', 'vocal': 'vocal'}[part_type]
  section = generator.arrange.form.ArrangementSection('', 0, len(prepared['staves'][0]['measures']), 0, role)
  with unittest.mock.patch('generator.arrange.form.full_arrangement_sections', return_value=[section]):
    if plan is not None:
      plans = {0: plan}
    elif part_type == 'acoustic':
      plans = generator.arrange.generate.plan_arrangement(prepared, rhythms)
    else:
      plans = {}
    return generator.arrange.generate._arrange_document(prepared, part_type, plans)