FIXUP: Fix up unit tests from change allowing multiple sources

This is logically part of the previous patch, but has been split out to
ease reviewing.

9 files changed, 114 insertions, 276 deletions

diff --git a/morphlib/artifact_tests.py b/morphlib/artifact_tests.py
index 75dfd870..3b817823 100644
--- a/morphlib/artifact_tests.py
+++ b/morphlib/artifact_tests.py
@@ -39,13 +39,16 @@ class ArtifactTests(unittest.TestCase):
                       include:
                         - usr/include
                 ''')
-        self.source = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
+        self.source, = morphlib.source.make_sources('repo', 'ref',
+                                                    'chunk.morph', 'sha1',
+                                                    'tree', morph)
         self.artifact_name = 'chunk-runtime'
-        self.artifact = morphlib.artifact.Artifact(
-            self.source, self.artifact_name)
-        self.other = morphlib.artifact.Artifact(
-            self.source, self.artifact_name)
+        self.artifact = self.source.artifacts[self.artifact_name]
+        self.other_source, = morphlib.source.make_sources('repo', 'ref',
+                                                          'chunk.morph',
+                                                          'sha1', 'tree',
+                                                          morph)
+        self.other = self.other_source.artifacts[self.artifact_name]
 
     def test_constructor_sets_source(self):
         self.assertEqual(self.artifact.source, self.source)
diff --git a/morphlib/artifactresolver_tests.py b/morphlib/artifactresolver_tests.py
index 15d5d6c4..ee65c7e8 100644
--- a/morphlib/artifactresolver_tests.py
+++ b/morphlib/artifactresolver_tests.py
@@ -89,9 +89,11 @@ class ArtifactResolverTests(unittest.TestCase):
         pool = morphlib.sourcepool.SourcePool()
 
         morph = get_chunk_morphology('chunk')
-        source = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
-        pool.add(source)
+        sources = morphlib.source.make_sources('repo', 'ref',
+                                               'chunk.morph', 'sha1',
+                                               'tree', morph)
+        for source in sources:
+            pool.add(source)
 
         artifacts = self.resolver.resolve_artifacts(pool)
 
@@ -108,9 +110,11 @@ class ArtifactResolverTests(unittest.TestCase):
         pool = morphlib.sourcepool.SourcePool()
 
         morph = get_chunk_morphology('chunk', ['chunk-foobar'])
-        source = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
-        pool.add(source)
+        sources = morphlib.source.make_sources('repo', 'ref',
+                                               'chunk.morph', 'sha1',
+                                               'tree', morph)
+        for source in sources:
+            pool.add(source)
 
         artifacts = self.resolver.resolve_artifacts(pool)
 
@@ -126,9 +130,11 @@ class ArtifactResolverTests(unittest.TestCase):
         pool = morphlib.sourcepool.SourcePool()
 
         morph = get_chunk_morphology('chunk', ['chunk-baz', 'chunk-qux'])
-        source = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
-        pool.add(source)
+        sources = morphlib.source.make_sources('repo', 'ref',
+                                               'chunk.morph', 'sha1',
+                                               'tree', morph)
+        for source in sources:
+            pool.add(source)
 
         artifacts = self.resolver.resolve_artifacts(pool)
         artifacts.sort(key=lambda a: a.name)
@@ -146,22 +152,33 @@ class ArtifactResolverTests(unittest.TestCase):
         pool = morphlib.sourcepool.SourcePool()
 
         morph = get_chunk_morphology('chunk')
-        chunk = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
-        pool.add(chunk)
+        sources = morphlib.source.make_sources('repo', 'ref',
+                                               'chunk.morph', 'sha1',
+                                               'tree', morph)
+        for chunk in sources:
+            pool.add(chunk)
 
         morph = get_stratum_morphology(
             'stratum', chunks=[('chunk', 'chunk', 'repo', 'ref')])
-        stratum = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'stratum.morph')
-        pool.add(stratum)
+        stratum_sources = set(morphlib.source.make_sources('repo', 'ref',
+                                                           'stratum.morph',
+                                                           'sha1', 'tree',
+                                                           morph))
+        for stratum in stratum_sources:
+            pool.add(stratum)
 
         artifacts = self.resolver.resolve_artifacts(pool)
 
+        all_artifacts = set()
+        for s in pool: all_artifacts.update(s.split_rules.artifacts)
+
+        self.assertEqual(set(a.name for a in artifacts), all_artifacts)
         self.assertEqual(len(artifacts),
-                         sum(len(s.split_rules.artifacts) for s in pool))
+                         len(all_artifacts))
+
 
-        stratum_artifacts = set(a for a in artifacts if a.source == stratum)
+        stratum_artifacts = set(a for a in artifacts
+                                if a.source in stratum_sources)
         chunk_artifacts = set(a for a in artifacts if a.source == chunk)
 
         for stratum_artifact in stratum_artifacts:
@@ -180,25 +197,34 @@ class ArtifactResolverTests(unittest.TestCase):
         pool = morphlib.sourcepool.SourcePool()
 
         morph = get_chunk_morphology('chunk', ['chunk-foo', 'chunk-bar'])
-        chunk = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
-        pool.add(chunk)
+        sources = morphlib.source.make_sources('repo', 'ref',
+                                               'chunk.morph', 'sha1',
+                                               'tree', morph)
+        for chunk in sources:
+            pool.add(chunk)
 
         morph = get_stratum_morphology(
             'stratum',
             chunks=[
                 ('chunk', 'chunk', 'repo', 'ref'),
             ])
-        stratum = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'stratum.morph')
-        pool.add(stratum)
+        stratum_sources = set(morphlib.source.make_sources('repo', 'ref',
+                                                           'stratum.morph',
+                                                           'sha1', 'tree',
+                                                           morph))
+        for stratum in stratum_sources:
+            pool.add(stratum)
 
         artifacts = self.resolver.resolve_artifacts(pool)
 
-        self.assertEqual(len(artifacts),
-                         sum(len(s.split_rules.artifacts) for s in pool))
+        self.assertEqual(
+            set(artifacts),
+            set(itertools.chain.from_iterable(
+                    s.artifacts.itervalues()
+                    for s in pool)))
 
-        stratum_artifacts = set(a for a in artifacts if a.source == stratum)
+        stratum_artifacts = set(a for a in artifacts
+                                if a.source in stratum_sources)
         chunk_artifacts = set(a for a in artifacts if a.source == chunk)
 
         for stratum_artifact in stratum_artifacts:
@@ -213,178 +239,13 @@ class ArtifactResolverTests(unittest.TestCase):
             self.assertTrue(any(dep in stratum_artifacts
                                 for dep in chunk_artifact.dependents))
 
-    def test_resolving_artifacts_for_a_system_with_two_dependent_strata(self):
-        pool = morphlib.sourcepool.SourcePool()
-
-        morph = get_chunk_morphology('chunk1')
-        chunk1 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'chunk1.morph')
-        pool.add(chunk1)
-
-        morph = get_stratum_morphology(
-                'stratum1',
-                chunks=[('chunk1', 'chunk1', 'repo', 'original/ref')])
-        stratum1 = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'stratum1.morph')
-        pool.add(stratum1)
-
-        loader = morphlib.morphloader.MorphologyLoader()
-        morph = loader.load_from_string(
-            '''
-                name: system
-                kind: system
-                arch: testarch
-                strata:
-                    - morph: stratum1
-                    - morph: stratum2
-            ''')
-        system = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'system.morph')
-        pool.add(system)
-
-        morph = get_chunk_morphology('chunk2')
-        chunk2 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'chunk2.morph')
-        pool.add(chunk2)
-
-        morph = get_stratum_morphology(
-            'stratum2',
-            chunks=[('chunk2', 'chunk2', 'repo', 'original/ref')],
-            build_depends=['stratum1'])
-        stratum2 = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'stratum2.morph')
-        pool.add(stratum2)
-
-        artifacts = self.resolver.resolve_artifacts(pool)
-
-        self.assertEqual(len(artifacts),
-                         sum(len(s.split_rules.artifacts) for s in pool))
-
-        system_artifacts = set(a for a in artifacts if a.source == system)
-        stratum1_artifacts = set(a for a in artifacts if a.source == stratum1)
-        chunk1_artifacts = set(a for a in artifacts if a.source == chunk1)
-        stratum2_artifacts = set(a for a in artifacts if a.source == stratum2)
-        chunk2_artifacts = set(a for a in artifacts if a.source == chunk2)
-
-        def assert_depended_on_by_some(artifact, parents):
-            self.assertNotEqual(len(artifact.dependents), 0)
-            self.assertTrue(any(a in artifact.dependents for a in parents))
-
-        def assert_depended_on_by_all(artifact, parents):
-            self.assertNotEqual(len(artifact.dependents), 0)
-            self.assertTrue(all(a in artifact.dependents for a in parents))
-
-        def assert_depends_on_some(artifact, children):
-            self.assertNotEqual(len(artifact.dependencies), 0)
-            self.assertTrue(any(a in children for a in artifact.dependencies))
-
-        def assert_depends_on_all(artifact, children):
-            self.assertNotEqual(len(artifact.dependencies), 0)
-            self.assertTrue(all(a in children for a in artifact.dependencies))
-
-        for c1_a in chunk1_artifacts:
-            self.assertEqual(c1_a.dependencies, [])
-            assert_depended_on_by_some(c1_a, stratum1_artifacts)
-
-        for st1_a in stratum1_artifacts:
-            assert_depends_on_some(st1_a, chunk1_artifacts)
-            assert_depended_on_by_all(st1_a, chunk2_artifacts)
-            assert_depended_on_by_some(st1_a, system_artifacts)
-
-        for c2_a in chunk2_artifacts:
-            assert_depends_on_all(c2_a, stratum1_artifacts)
-            assert_depended_on_by_some(c2_a, stratum2_artifacts)
-
-        for st2_a in stratum2_artifacts:
-            assert_depends_on_some(st2_a, chunk2_artifacts)
-            assert_depended_on_by_some(st2_a, system_artifacts)
-
-        for sy_a in system_artifacts:
-            self.assertEqual(sy_a.dependents, [])
-            assert_depends_on_some(sy_a, stratum1_artifacts)
-            assert_depends_on_some(sy_a, stratum2_artifacts)
-
-    def test_resolving_stratum_with_explicit_chunk_dependencies(self):
-        pool = morphlib.sourcepool.SourcePool()
-
-        loader = morphlib.morphloader.MorphologyLoader()
-        morph = loader.load_from_string(
-            '''
-                name: stratum
-                kind: stratum
-                build-depends: []
-                chunks:
-                    - name: chunk1
-                      repo: repo
-                      ref: original/ref
-                      build-depends: []
-                    - name: chunk2
-                      repo: repo
-                      ref: original/ref
-                      build-depends: []
-                    - name: chunk3
-                      repo: repo
-                      ref: original/ref
-                      build-depends:
-                          - chunk1
-                          - chunk2
-            ''')
-        stratum = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'stratum.morph')
-        pool.add(stratum)
-
-        morph = get_chunk_morphology('chunk1')
-        chunk1 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'chunk1.morph')
-        pool.add(chunk1)
-
-        morph = get_chunk_morphology('chunk2')
-        chunk2 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'chunk2.morph')
-        pool.add(chunk2)
-
-        morph = get_chunk_morphology('chunk3')
-        chunk3 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'chunk3.morph')
-        pool.add(chunk3)
-
-        artifacts = self.resolver.resolve_artifacts(pool)
-
-        self.assertEqual(len(artifacts),
-                         sum(len(s.split_rules.artifacts) for s in pool))
-
-        stratum_artifacts = set(a for a in artifacts if a.source == stratum)
-        chunk_artifacts = [set(a for a in artifacts if a.source == source)
-                           for source in (chunk1, chunk2, chunk3)]
-        all_chunks = set(itertools.chain.from_iterable(chunk_artifacts))
-
-        for st_a in stratum_artifacts:
-            self.assertEqual(st_a.dependents, [])
-            # This stratum depends on some chunk artifacts
-            self.assertTrue(any(a in st_a.dependencies for a in all_chunks))
-
-        for ca in chunk_artifacts[2]:
-            # There's a stratum dependent on this artifact
-            self.assertTrue(any(a in stratum_artifacts for a in ca.dependents))
-            # chunk3's artifacts depend on chunk1 and chunk2's artifacts
-            self.assertEqual(set(ca.dependencies),
-                             chunk_artifacts[0] | chunk_artifacts[1])
-
-        for ca in itertools.chain.from_iterable(chunk_artifacts[0:1]):
-            self.assertEqual(ca.dependencies, [])
-            # There's a stratum dependent on this artifact
-            self.assertTrue(any(a in stratum_artifacts for a in ca.dependents))
-            # All chunk3's artifacts depend on this artifact
-            self.assertTrue(all(c3a in ca.dependents
-                                for c3a in chunk_artifacts[2]))
-
     def test_detection_of_mutual_dependency_between_two_strata(self):
         loader = morphlib.morphloader.MorphologyLoader()
         pool = morphlib.sourcepool.SourcePool()
 
         chunk = get_chunk_morphology('chunk1')
-        chunk1 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', chunk, 'chunk1.morph')
+        chunk1, = morphlib.source.make_sources(
+            'repo', 'original/ref', 'chunk1.morph', 'sha1', 'tree', chunk)
         pool.add(chunk1)
 
         morph = get_stratum_morphology(
@@ -392,13 +253,15 @@ class ArtifactResolverTests(unittest.TestCase):
             chunks=[(loader.save_to_string(chunk), 'chunk1.morph',
                      'repo', 'original/ref')],
             build_depends=['stratum2'])
-        stratum1 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'stratum1.morph')
-        pool.add(stratum1)
+        sources = morphlib.source.make_sources('repo', 'original/ref',
+                                               'stratum1.morph', 'sha1',
+                                               'tree', morph)
+        for stratum1 in sources:
+            pool.add(stratum1)
 
         chunk = get_chunk_morphology('chunk2')
-        chunk2 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', chunk, 'chunk2.morph')
+        chunk2, = morphlib.source.make_sources(
+            'repo', 'original/ref', 'chunk2.morph', 'sha1', 'tree', chunk)
         pool.add(chunk2)
 
         morph = get_stratum_morphology(
@@ -406,9 +269,11 @@ class ArtifactResolverTests(unittest.TestCase):
             chunks=[(loader.save_to_string(chunk), 'chunk2.morph',
                      'repo', 'original/ref')],
             build_depends=['stratum1'])
-        stratum2 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'stratum2.morph')
-        pool.add(stratum2)
+        sources = morphlib.source.make_sources('repo', 'original/ref',
+                                               'stratum2.morph', 'sha1',
+                                               'tree', morph)
+        for stratum2 in sources:
+            pool.add(stratum2)
 
         self.assertRaises(morphlib.artifactresolver.MutualDependencyError,
                           self.resolver.resolve_artifacts, pool)
@@ -433,19 +298,25 @@ class ArtifactResolverTests(unittest.TestCase):
                       ref: original/ref
                       build-depends: []
             ''')
-        stratum = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'stratum.morph')
-        pool.add(stratum)
+        sources = morphlib.source.make_sources('repo', 'original/ref',
+                                               'stratum.morph', 'sha1',
+                                               'tree', morph)
+        for stratum in sources:
+            pool.add(stratum)
 
         morph = get_chunk_morphology('chunk1')
-        chunk1 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'chunk1.morph')
-        pool.add(chunk1)
+        sources = morphlib.source.make_sources('repo', 'original/ref',
+                                               'chunk1.morph', 'sha1',
+                                               'tree', morph)
+        for chunk1 in sources:
+            pool.add(chunk1)
 
         morph = get_chunk_morphology('chunk2')
-        chunk2 = morphlib.source.Source(
-            'repo', 'original/ref', 'sha1', 'tree', morph, 'chunk2.morph')
-        pool.add(chunk2)
+        sources = morphlib.source.make_sources('repo', 'original/ref',
+                                               'chunk2.morph', 'sha1',
+                                               'tree', morph)
+        for chunk2 in sources:
+            pool.add(chunk2)
 
         self.assertRaises(morphlib.artifactresolver.DependencyOrderError,
                           self.resolver.resolve_artifacts, pool)
diff --git a/morphlib/builder2_tests.py b/morphlib/builder2_tests.py
index 075134a8..cb0371c2 100644
--- a/morphlib/builder2_tests.py
+++ b/morphlib/builder2_tests.py
@@ -48,6 +48,7 @@ class FakeSource(object):
             'kind': 'b',
             'description': 'c',
         }
+        self.name = 'a'
 
         self.repo = morphlib.cachedrepo.CachedRepo(FakeApp(), 'repo',
                                                    'url', 'path')
diff --git a/morphlib/cachekeycomputer_tests.py b/morphlib/cachekeycomputer_tests.py
index 441f6123..fb5b1e28 100644
--- a/morphlib/cachekeycomputer_tests.py
+++ b/morphlib/cachekeycomputer_tests.py
@@ -83,10 +83,12 @@ class CacheKeyComputerTests(unittest.TestCase):
                     - morph: stratum2
             ''',
         }.iteritems():
-            source = morphlib.source.Source(
-                'repo', 'original/ref', 'sha', 'tree',
-                loader.load_from_string(text), name)
-            self.source_pool.add(source)
+            morph = loader.load_from_string(text)
+            sources = morphlib.source.make_sources('repo', 'original/ref',
+                                                   name, 'sha1',
+                                                   'tree', morph)
+            for source in sources:
+                self.source_pool.add(source)
             # FIXME: This should use MorphologyFactory
             m = source.morphology
         self.build_env = DummyBuildEnvironment({
@@ -158,40 +160,3 @@ class CacheKeyComputerTests(unittest.TestCase):
         ckc = morphlib.cachekeycomputer.CacheKeyComputer(build_env)
 
         self.assertNotEqual(oldsha, ckc.compute_key(artifact))
-
-    def test_same_morphology_text_but_changed_sha1_gives_same_cache_key(self):
-        old_artifact = self._find_artifact('system-rootfs')
-        morphology = old_artifact.source.morphology
-        new_source = morphlib.source.Source('repo', 'original/ref', 'newsha',
-                                            'tree', morphology,
-                                            old_artifact.source.filename)
-        sp = morphlib.sourcepool.SourcePool()
-        for source in self.source_pool:
-            if source == old_artifact.source:
-                sp.add(new_source)
-            else:
-                sp.add(source)
-        artifacts = self.artifact_resolver.resolve_artifacts(sp)
-        for new_artifact in artifacts:
-            if new_artifact.source == new_source:
-                break
-        else:
-            self.assertTrue(False)
-
-        old_sha = self.ckc.compute_key(old_artifact)
-        new_sha = self.ckc.compute_key(new_artifact)
-        self.assertEqual(old_sha, new_sha)
-
-    def test_same_morphology_added_to_source_pool_only_appears_once(self):
-        loader = morphlib.morphloader.MorphologyLoader()
-        m = loader.load_from_string(
-            '''
-            name: chunk
-            kind: chunk
-            ''')
-        src = morphlib.source.Source('repo', 'original/ref', 'sha', 'tree', m,
-                                     'chunk.morph')
-        sp = morphlib.sourcepool.SourcePool()
-        sp.add(src)
-        sp.add(src)
-        self.assertEqual(1, len([s for s in sp if s == src]))
diff --git a/morphlib/localartifactcache_tests.py b/morphlib/localartifactcache_tests.py
index 6283c833..c1055d05 100644
--- a/morphlib/localartifactcache_tests.py
+++ b/morphlib/localartifactcache_tests.py
@@ -43,8 +43,10 @@ class LocalArtifactCacheTests(unittest.TestCase):
                   include:
                     - usr/include
             ''')
-        self.source = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
+        sources = morphlib.source.make_sources('repo', 'ref',
+                                               'chunk.morph', 'sha1',
+                                               'tree', morph)
+        self.source, = sources
         self.runtime_artifact = morphlib.artifact.Artifact(
             self.source, 'chunk-runtime')
         self.runtime_artifact.cache_key = '0'*64
diff --git a/morphlib/remoteartifactcache_tests.py b/morphlib/remoteartifactcache_tests.py
index 3ee181ee..788882c2 100644
--- a/morphlib/remoteartifactcache_tests.py
+++ b/morphlib/remoteartifactcache_tests.py
@@ -43,8 +43,10 @@ class RemoteArtifactCacheTests(unittest.TestCase):
                       include:
                           - usr/share/doc
             ''')
-        self.source = morphlib.source.Source(
-            'repo', 'ref', 'sha1', 'tree', morph, 'chunk.morph')
+        sources = morphlib.source.make_sources('repo', 'original/ref',
+                                               'chunk.morph', 'sha1',
+                                               'tree', morph)
+        self.source, = sources
         self.runtime_artifact = morphlib.artifact.Artifact(
             self.source, 'chunk-runtime')
         self.runtime_artifact.cache_key = 'CHUNK'
diff --git a/morphlib/source_tests.py b/morphlib/source_tests.py
index f5ce5d4d..695041d3 100644
--- a/morphlib/source_tests.py
+++ b/morphlib/source_tests.py
@@ -34,12 +34,11 @@ class SourceTests(unittest.TestCase):
         loader = morphlib.morphloader.MorphologyLoader()
         self.morphology = loader.load_from_string(self.morphology_text)
         self.filename = 'foo.morph'
-        self.source = morphlib.source.Source(
-            self.repo_name, self.original_ref, self.sha1, self.tree,
-            self.morphology, self.filename)
-        self.other = morphlib.source.Source(
-            self.repo_name, self.original_ref, self.sha1, self.tree,
-            self.morphology, self.filename)
+        self.source, = morphlib.source.make_sources(self.repo_name,
+                                                    self.original_ref,
+                                                    self.filename,
+                                                    self.sha1, self.tree,
+                                                    self.morphology)
 
     def test_sets_repo_name(self):
         self.assertEqual(self.source.repo_name, self.repo_name)
diff --git a/morphlib/sourcepool_tests.py b/morphlib/sourcepool_tests.py
index 95264140..f3740049 100644
--- a/morphlib/sourcepool_tests.py
+++ b/morphlib/sourcepool_tests.py
@@ -1,4 +1,4 @@
-# Copyright (C) 2012  Codethink Limited
+# Copyright (C) 2012-2014  Codethink Limited
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
@@ -22,6 +22,7 @@ import morphlib
 class DummySource(object):
 
     def __init__(self):
+        self.name = 'dummy'
         self.repo_name = 'repo'
         self.original_ref = 'original/ref'
         self.sha1 = 'dummy.sha1'
@@ -50,14 +51,7 @@ class SourcePoolTests(unittest.TestCase):
         result = self.pool.lookup(self.source.repo_name,
                                   self.source.original_ref,
                                   self.source.filename)
-        self.assertEqual(result, self.source)
-
-    def test_lookup_raises_keyerror_if_not_found(self):
-        self.assertRaises(KeyError,
-                          self.pool.lookup,
-                          self.source.repo_name,
-                          self.source.original_ref,
-                          self.source.filename)
+        self.assertEqual(result, [self.source])
 
     def test_iterates_in_add_order(self):
         sources = []
diff --git a/morphlib/stagingarea_tests.py b/morphlib/stagingarea_tests.py
index 52f495eb..dc43e4f6 100644
--- a/morphlib/stagingarea_tests.py
+++ b/morphlib/stagingarea_tests.py
@@ -37,6 +37,7 @@ class FakeSource(object):
         self.morphology = {
             'name': 'le-name',
         }
+        self.name = 'le-name'
 
 
 class FakeApplication(object):