Restart codereview issue 6792042

buildbot/cbuildbot_stages_unittest.py

 #!/usr/bin/python
 
 # Copyright (c) 2011 The Chromium OS Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 """Unittests for build stages."""
 
 import mox
 import os
@@ skipping 604 matching lines @@
      return stages.BuildTargetStage(self.bot_id,
                                     self.options,
                                     self.build_config)
 
   def testAllConditionalPaths(self):
     """Enable all paths to get line coverage."""
     self.build_config['vm_tests'] = True
     self.options.tests = True
     self.build_config['build_type'] = 'full'
     self.build_config['usepkg'] = True
+    self.build_config['useflags'] = ['ALPHA', 'BRAVO', 'CHARLIE']
+    proper_env = {'USE' : ' '.join(self.build_config['useflags'])}
 
     stages.BuilderStage._GetPortageEnvVar('FULL_BINHOST').AndReturn('new.com')
     stages.BuilderStage.old_binhost = 'old.com'
 
     commands.Build(
-        self.build_root, True,
-        build_autotest=True, usepkg=True)
+        self.build_root, True, build_autotest=True, usepkg=True,
+        extra_env=proper_env)
 
     commands.UploadPrebuilts(
         self.build_root, self.build_config['board'],
         self.build_config['rev_overlays'], [],
         self.build_config['build_type'], False)
 
-    commands.BuildImage(self.build_root)
-    commands.BuildVMImageForTesting(self.build_root)
+    commands.BuildImage(self.build_root, extra_env=proper_env)
+    commands.BuildVMImageForTesting(self.build_root, extra_env=proper_env)
 
     self.mox.ReplayAll()
     self.RunStage()
     self.mox.VerifyAll()
 
   def testFalseBuildArgs1(self):
     """Make sure our logic for Build arguments can toggle to false."""
     stages.BuilderStage._GetPortageEnvVar('FULL_BINHOST').AndReturn('new.com')
     stages.BuilderStage.old_binhost = 'new.com'
+    self.build_config['useflags'] = None
 
-    commands.Build(self.build_root, False, build_autotest=False, usepkg=False)
-    commands.BuildImage(self.build_root)
+    commands.Build(self.build_root, False, build_autotest=False, usepkg=False,
+        extra_env=None)
+    commands.BuildImage(self.build_root, extra_env=None)
 
     self.mox.ReplayAll()
     self.RunStage()
     self.mox.VerifyAll()
 
   def testFalseBuildArgs2(self):
     """Verify emptytree flag is false when there is no old binhost."""
     stages.BuilderStage._GetPortageEnvVar('FULL_BINHOST').AndReturn('new.com')
     stages.BuilderStage.old_binhost = False
+    self.build_config['useflags'] = ['BAKER']
+    proper_env = {'USE' : ' '.join(self.build_config['useflags'])}
 
     commands.Build(
         self.build_root, False,
-        build_autotest=False, usepkg=False)
+        build_autotest=False, usepkg=False, extra_env=proper_env)
 
-    commands.BuildImage(self.build_root)
+    commands.BuildImage(self.build_root, extra_env=proper_env)
 
     self.mox.ReplayAll()
     self.RunStage()
     self.mox.VerifyAll()
 
 
 class ArchiveStageTest(AbstractStageTest):
 
   def setUp(self):
     mox.MoxTestBase.setUp(self)
@@ skipping 248 matching lines @@
     actualResults = stages.BuilderStage.Results.Get()
 
     # Break out the asserts to be per item to make debugging easier
     self.assertEqual(len(expectedResults), len(actualResults))
     for i in xrange(len(expectedResults)):
       self.assertEqual(expectedResults[i], actualResults[i])
 
 
 if __name__ == '__main__':
   unittest.main()