From d9619e993a9b84d8afa3f81854169d0099778c3c Mon Sep 17 00:00:00 2001
From: Joakim Soderlund <joakim.soderlund@gmail.com>
Date: Sat, 4 Nov 2017 20:41:22 +0100
Subject: [PATCH] Add stamper to update task

---
 fimfarchive/tasks/update.py |  8 ++++++++
 tests/tasks/test_update.py  | 15 ++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/fimfarchive/tasks/update.py b/fimfarchive/tasks/update.py
index a93b570..8fa1402 100644
--- a/fimfarchive/tasks/update.py
+++ b/fimfarchive/tasks/update.py
@@ -32,6 +32,7 @@ from fimfarchive.flavors import DataFormat, UpdateStatus
 from fimfarchive.mappers import StoryPathMapper
 from fimfarchive.selectors import Selector, UpdateSelector
 from fimfarchive.signals import Signal, SignalSender
+from fimfarchive.stampers import Stamper, UpdateStamper
 from fimfarchive.stories import Story
 from fimfarchive.utils import PersistedDict
 from fimfarchive.writers import DirectoryWriter
@@ -75,6 +76,7 @@ class UpdateTask(SignalSender):
             fimfarchive: Fetcher,
             fimfiction: Fetcher,
             selector: Selector = None,
+            stamper: Stamper = None,
             workdir: str = DEFAULT_WORKDIR,
             retries: int = DEFAULT_RETRIES,
             skips: int = DEFAULT_SKIPS,
@@ -86,6 +88,7 @@ class UpdateTask(SignalSender):
             fimfarchive: Fetcher for the old release.
             fimfiction: Fetcher for the new release.
             selector: Selector for which story to save.
+            stamper: Stamper for adding story information.
             workdir: Path for storage of state and stories.
             retries: Number of retries before giving up.
             skips: Number of skips before giving up.
@@ -95,9 +98,13 @@ class UpdateTask(SignalSender):
         if selector is None:
             selector = UpdateSelector()
 
+        if stamper is None:
+            stamper = UpdateStamper()
+
         self.fimfarchive = fimfarchive
         self.fimfiction = fimfiction
         self.select = selector
+        self.stamp = stamper
         self.workdir = workdir
         self.retries = retries
         self.skips = skips
@@ -176,6 +183,7 @@ class UpdateTask(SignalSender):
             selected = selected.merge(meta=new.meta)
 
         if selected:
+            self.stamp(selected)
             self.write(selected)
         elif new:
             self.skip_writer.write(new)
diff --git a/tests/tasks/test_update.py b/tests/tasks/test_update.py
index 40a367f..0b3e419 100644
--- a/tests/tasks/test_update.py
+++ b/tests/tasks/test_update.py
@@ -105,7 +105,14 @@ class TestUpdateTask:
         return UpdateSelector()
 
     @pytest.fixture
-    def task(self, fimfarchive, fimfiction, selector, tmpdir):
+    def stamper(self):
+        """
+        Returns a stamper mock.
+        """
+        return MagicMock()
+
+    @pytest.fixture
+    def task(self, fimfarchive, fimfiction, selector, stamper, tmpdir):
         """
         Returns an `UpdateTask` instance.
         """
@@ -113,6 +120,7 @@ class TestUpdateTask:
             fimfiction=fimfiction,
             fimfarchive=fimfarchive,
             selector=selector,
+            stamper=stamper,
             workdir=str(tmpdir),
             retries=2,
             skips=2,
@@ -142,6 +150,7 @@ class TestUpdateTask:
         )
 
         self.verify_run(task, delays)
+        task.stamp.assert_called_once_with(target)
         task.write.assert_called_once_with(target)
         assert status in target.flavors
         assert task.state['key'] == 4
@@ -150,6 +159,7 @@ class TestUpdateTask:
         """
         Runs the task and verifies an empty fetch.
         """
+        task.write = MagicMock()
         task.skip_writer.write = MagicMock()
 
         target = fetcher.add(key=1, date=1)
@@ -161,6 +171,8 @@ class TestUpdateTask:
         )
 
         self.verify_run(task, delays)
+        task.stamp.assert_not_called()
+        task.write.assert_not_called()
         task.skip_writer.write.assert_called_once_with(target)
 
     def verify_failure(self, task, fetcher):
@@ -176,6 +188,7 @@ class TestUpdateTask:
         )
 
         self.verify_run(task, delays)
+        task.stamp.assert_not_called()
         task.write.assert_not_called()
 
     def test_created_story(self, task, fimfiction):