def handle(self, **options):
importer = YNRBallotImporter(
stdout=self.stdout,
current_only=options["current"],
force_metadata=options["force_metadata"],
)
importer.do_import()
def setUp(self):
self.importer = YNRBallotImporter()
election = ElectionFactory(name="Adur local election")
self.post = PostFactory(label="Adur local election")
self.post.elections.add(election)
self.orphan_post = PostFactory(label="Adur local election",
ynr_id="foo")
def handle(self, **options):
importer = YNRBallotImporter(
stdout=self.stdout,
current_only=options["current"],
force_metadata=options["force_metadata"],
force_current_metadata=options["force_current_metadata"],
)
importer.do_import()
self.populate_any_non_by_elections_field()
def test_set_division_type_force_no_ee_data(self, ballot):
"""
Test that force_update ensures update attempt tries, but doesnt change
when no data
"""
importer = YNRBallotImporter(force_update=True)
assert importer.set_division_type(ballot=ballot) is None
assert ballot.post.division_type == "DIW"
EEHelper.get_data.assert_called_once_with(ballot.ballot_paper_id)
ballot.post.save.assert_not_called()
def test_set_division_type_unchanged(self, ballot):
"""
Test that when division_type is set and not a force update, nothing
happens
"""
importer = YNRBallotImporter(force_update=False)
assert importer.set_division_type(ballot=ballot) is None
assert ballot.post.division_type == "DIW"
EEHelper.get_data.assert_not_called()
ballot.post.save.assert_not_called()
def test_set_division_type_changed(self, ballot, mocker):
"""
Test that division_type gets updated
"""
importer = YNRBallotImporter(force_update=True)
division = {"division": {"division_type": "NEW"}}
mocker.patch.object(EEHelper, "get_data", return_value=division)
mocker.patch.object(ballot.post, "full_clean")
assert importer.set_division_type(ballot=ballot) is None
assert ballot.post.division_type == "NEW"
EEHelper.get_data.assert_called_once_with(ballot.ballot_paper_id)
ballot.post.full_clean.assert_called_once()
ballot.post.save.assert_called_once()
def handle(self, **options):
self.options = options
self.dirpath = tempfile.mkdtemp()
self.ballot_importer = YNRBallotImporter(stdout=self.stdout)
try:
self.past_time_str = Person.objects.latest().last_updated
except Person.DoesNotExist:
# In case this is the first run
self.past_time_str = "1800-01-01"
if self.options["since"]:
self.past_time_str = self.options["since"]
try:
self.download_pages()
self.add_to_db()
finally:
self.delete_merged_people()
shutil.rmtree(self.dirpath)
class TestYNRBallotImporter(TestCase):
def setUp(self):
self.importer = YNRBallotImporter()
election = ElectionFactory(name="Adur local election")
self.post = PostFactory(label="Adur local election")
self.post.elections.add(election)
self.orphan_post = PostFactory(label="Adur local election",
ynr_id="foo")
def test_delete_orphan_posts(self):
deleted_posts, _ = self.importer.delete_orphan_posts()
query_set = Post.objects.all()
self.assertEqual(deleted_posts, 1)
self.assertEqual(query_set.count(), 1)
self.assertIn(self.post, query_set)
self.assertNotIn(self.orphan_post, query_set)
class Command(BaseCommand):
def add_arguments(self, parser):
parser.add_argument(
"--recent",
action="store_true",
dest="recent",
default=False,
help="Import changes in the last `n` minutes",
)
parser.add_argument(
"--since",
action="store",
dest="since",
type=self.valid_date,
help="Import changes since [datetime]",
)
parser.add_argument(
"--update-info-only",
action="store_true",
help="Only update person info, not posts",
)
def valid_date(self, value):
return parse(value)
def handle(self, **options):
self.options = options
self.dirpath = tempfile.mkdtemp()
self.ballot_importer = YNRBallotImporter(stdout=self.stdout)
try:
last_updated = Person.objects.latest().last_updated
self.past_time_str = last_updated - timedelta(hours=1)
except Person.DoesNotExist:
# In case this is the first run
self.past_time_str = "1800-01-01"
if self.options["since"]:
self.past_time_str = self.options["since"]
self.past_time_str = str(self.past_time_str)
try:
self.download_pages()
self.add_to_db()
finally:
self.delete_merged_people()
shutil.rmtree(self.dirpath)
def add_to_db(self):
self.existing_people = set(Person.objects.values_list("pk", flat=True))
self.seen_people = set()
files = [f for f in os.listdir(self.dirpath) if f.endswith(".json")]
files = sorted(files,
key=lambda k: int(k.split("-")[-1].split(".")[0]))
for file in files:
self.stdout.write("Importing {}".format(file))
with open(os.path.join(self.dirpath, file), "r") as f:
results = json.loads(f.read())
self.add_people(
results, update_info_only=self.options["update_info_only"])
should_clean_up = not any([
self.options["recent"],
self.options["update_info_only"],
self.options["since"],
])
if should_clean_up:
deleted_ids = self.existing_people.difference(self.seen_people)
Person.objects.filter(ynr_id__in=deleted_ids).delete()
def save_page(self, url, page):
# get the file name from the page number
if "cached-api" in url:
filename = url.split("/")[-1]
else:
if "page=" in url:
page_number = url.split("page=")[1].split("&")[0]
else:
page_number = 1
filename = "page-{}.json".format(page_number)
file_path = os.path.join(self.dirpath, filename)
# Save the page
with open(file_path, "w") as f:
f.write(page)
def download_pages(self):
params = {"page_size": "200"}
if self.options["recent"] or self.options["since"]:
params["updated_gte"] = self.past_time_str
next_page = settings.YNR_BASE + "/api/next/people/?{}".format(
urlencode(params))
else:
next_page = (settings.YNR_BASE +
"/media/cached-api/latest/people-000001.json")
while next_page:
self.stdout.write("Downloading {}".format(next_page))
req = requests.get(next_page)
req.raise_for_status()
page = req.text
results = req.json()
self.save_page(next_page, page)
next_page = results.get("next")
@transaction.atomic
def add_people(self, results, update_info_only=False):
for person in results["results"]:
with show_data_on_error("Person {}".format(person["id"]), person):
person_obj = Person.objects.update_or_create_from_ynr(
person, update_info_only=update_info_only)
if self.options["recent"]:
self.update_candidacies(person, person_obj)
if person["candidacies"]:
self.seen_people.add(person_obj.pk)
def update_candidacies(self, person, person_obj):
"""
Compare the API dict to the candidacies we know about
"""
known_candidacies = set(person_obj.personpost_set.all().values_list(
"post_election__ballot_paper_id", flat=True))
remote_candidacies = set(
[c["ballot"]["ballot_paper_id"] for c in person["candidacies"]])
if remote_candidacies != known_candidacies:
# Delete any old candidacies
deleted_candidacies = known_candidacies - remote_candidacies
person_obj.personpost_set.filter(
post_election__ballot_paper_id__in=deleted_candidacies).delete(
)
added_candidacies = remote_candidacies - known_candidacies
for candidacy in added_candidacies:
try:
ballot = PostElection.objects.get(
ballot_paper_id=candidacy)
except PostElection.DoesNotExist:
# This might be because we've not run import_ballots
# recently enough. Let's import just the ballots for this
# date
self.import_ballots_for_date(candidacy.split(".")[-1])
ballot = PostElection.objects.get(
ballot_paper_id=candidacy)
for candidacy_dict in person["candidacies"]:
if candidacy_dict["ballot"][
"ballot_paper_id"] == candidacy:
candidacy_dict_for_ballot = candidacy_dict
person_obj.personpost_set.update_or_create(
post_election=ballot,
post=ballot.post,
election=ballot.election,
party_id=candidacy_dict_for_ballot["party"]["legacy_slug"],
list_position=candidacy_dict_for_ballot[
"party_list_position"],
elected=candidacy_dict_for_ballot["elected"],
)
def import_ballots_for_date(self, date):
self.ballot_importer.do_import(params={"election_date": date})
def delete_merged_people(self):
url = (
settings.YNR_BASE +
"/api/next/person_redirects/?page_size=200&updated_gte={}".format(
self.past_time_str))
merged_ids = []
while url:
req = requests.get(url)
page = req.json()
for result in page["results"]:
merged_ids.append(result["old_person_id"])
url = page.get("next")
Person.objects.filter(ynr_id__in=merged_ids).delete()