Morph build c3874f415dc6448ca28d9a01edab0948baserock/6453f312359f4317803ef7f14b58d21f/d675b946df4f456693ed211dcd2ec95e

System branch: master

1 files changed, 23 insertions, 4 deletions

diff --git a/morphlib/stagingarea_tests.py b/morphlib/stagingarea_tests.py
index ffdf5eaa..3d378573 100644
--- a/morphlib/stagingarea_tests.py
+++ b/morphlib/stagingarea_tests.py
@@ -46,6 +46,25 @@ class FakeArtifact(object):
         self.source = FakeSource()
 
 
+class FakeArtifactCache(object):
+
+    def __init__(self, tempdir):
+        self.tempdir = tempdir
+
+    def create_chunk(self, chunkdir):
+        if not chunkdir:
+            chunkdir = os.path.join(self.tempdir, 'chunk')
+        if not os.path.exists(chunkdir):
+            os.mkdir(chunkdir)
+        with open(os.path.join(chunkdir, 'file.txt'), 'w'):
+            pass
+
+        return chunkdir
+
+    def get(self, artifact, directory=None):
+        return self.create_chunk(directory)
+
+
 class FakeApplication(object):
 
     def __init__(self, cachedir, tempdir):
@@ -141,14 +160,14 @@ class StagingAreaTests(unittest.TestCase):
 
     def test_installs_artifact(self):
         artifact = FakeArtifact()
-        chunkdir = self.create_chunk()
-        self.sa.install_artifact(artifact, chunkdir)
+        artifact_cache = FakeArtifactCache(self.tempdir)
+        self.sa.install_artifact(artifact_cache, artifact)
         self.assertEqual(self.list_tree(self.staging), ['/', '/file.txt'])
 
     def test_removes_everything(self):
         artifact = FakeArtifact()
-        chunkdir = self.create_chunk()
-        self.sa.install_artifact(artifact, chunkdir)
+        artifact_cache = FakeArtifactCache(self.tempdir)
+        self.sa.install_artifact(artifact_cache, artifact)
         self.sa.remove()
         self.assertFalse(os.path.exists(self.staging))