From 5e1f81cc8e5c839b7cae9b3c1e56cbd5c9cebba7 Mon Sep 17 00:00:00 2001
From: Jan Kaluza <jkaluza@redhat.com>
Date: Dec 04 2017 12:47:48 +0000
Subject: Merge main ODCS compose with packages from advisory with Pulp compose.


---

diff --git a/freshmaker/handlers/__init__.py b/freshmaker/handlers/__init__.py
index 53ceac6..5590ed6 100644
--- a/freshmaker/handlers/__init__.py
+++ b/freshmaker/handlers/__init__.py
@@ -34,10 +34,13 @@ from freshmaker.types import ArtifactType, EventState
 from freshmaker.models import ArtifactBuild, Event
 from freshmaker.utils import krb_context, get_rebuilt_nvr
 from freshmaker.errors import UnprocessableEntity, ProgrammingError
+from freshmaker.consumer import work_queue_put
+from freshmaker.events import ODCSComposeStateChangeEvent
 
 from freshmaker.odcsclient import ODCS
 from freshmaker.odcsclient import AuthMech
 from freshmaker.odcsclient import COMPOSE_STATES
+from freshmaker.odcsclient import PUNGI_SOURCE_TYPE_NAMES
 
 
 def fail_event_on_handler_exception(func):
@@ -112,6 +115,9 @@ class BaseHandler(object):
     """
     __metaclass__ = abc.ABCMeta
 
+    # Used to generate incremental compose id in dry run mode.
+    _FAKE_COMPOSE_ID = 1
+
     def __init__(self):
         self._db_event_id = None
         self._db_artifact_build_id = None
@@ -260,6 +266,90 @@ class BaseHandler(object):
             raise UnprocessableEntity(err_msg)
         return in_whitelist
 
+    def odcs_get_compose(self, compose_id):
+        """
+        Returns the information from the ODCS server about compose with id
+        `compose_id`. In DRY_RUN mode, returns fake compose information
+        without contacting the ODCS server.
+        """
+        if conf.dry_run:
+            compose = {}
+            compose['id'] = compose_id
+            compose['result_repofile'] = "http://localhost/%d.repo" % (
+                compose['id'])
+            compose['state'] = COMPOSE_STATES['done']
+            return compose
+
+        odcs = ODCS(conf.odcs_server_url, auth_mech=AuthMech.Kerberos,
+                    verify_ssl=conf.odcs_verify_ssl)
+        with krb_context():
+            return odcs.get_compose(compose_id)
+
+    def _fake_odcs_new_compose(self, source_type, source, packages=None):
+        """
+        Fake ODCS.new_compose method used dry run mode.
+
+        Logs the arguments and emits ODCSComposeStateChangeEvent of
+        "done" state.
+
+        :rtype: dict
+        :return: Fake odcs.new_compose dict.
+        """
+        log.info("DRY RUN: Calling fake odcs.new_compose with args: %r",
+                 (source_type, source, packages))
+
+        # Generate the new_compose dict.
+        BaseHandler._FAKE_COMPOSE_ID += 1
+        new_compose = {}
+        new_compose['id'] = BaseHandler._FAKE_COMPOSE_ID
+        new_compose['result_repofile'] = "http://localhost/%d.repo" % (
+            new_compose['id'])
+        new_compose['source_type'] = PUNGI_SOURCE_TYPE_NAMES[source_type]
+        new_compose['source'] = source
+        new_compose['state'] = COMPOSE_STATES['done']
+
+        # Generate and inject the ODCSComposeStateChangeEvent event.
+        event = ODCSComposeStateChangeEvent(
+            "fake_odcs_compose_msg", new_compose)
+        log.info("Injecting fake event: %r", event)
+        work_queue_put(event)
+
+        return new_compose
+
+    def odcs_new_compose(self, source_type, source, packages=None, wait=False,
+                         result_repo_name=None):
+        """
+        Submites new ODCS request for a compose.
+
+        :param str source_type: ODCS compose source type.
+        :param str source: ODCS compose source.
+        :param list packages: List of packages to include in compose or None.
+        :param bool wait: When True, this method will block until the compose
+            is done.
+        :param str result_repo_name: Name of the resulting ODCS compose repo.
+        :rtype: dict
+        :return: Dict returned by ODCS.new_compose(...).
+        """
+        log.info('Generate new compose for rebuild: '
+                 'source_type: %s, source: %s, packages: %s',
+                 source_type, source, packages)
+
+        odcs = ODCS(conf.odcs_server_url, auth_mech=AuthMech.Kerberos,
+                    verify_ssl=conf.odcs_verify_ssl)
+        if not conf.dry_run:
+            with krb_context():
+                new_compose = odcs.new_compose(
+                    source, source_type, packages=packages,
+                    sigkeys=conf.odcs_sigkeys, flags=["no_deps"])
+                if wait:
+                    new_compose = odcs.wait_for_compose(new_compose['id'])
+        else:
+            new_compose = self._fake_odcs_new_compose(
+                source_type, source, packages=packages,
+                result_repo_name=result_repo_name)
+
+        return new_compose
+
 
 class ContainerBuildHandler(BaseHandler):
     """Handler for building containers"""
@@ -343,25 +433,6 @@ class ContainerBuildHandler(BaseHandler):
             scm_url, branch, target, repo_urls=repo_urls,
             isolated=True, release=release, koji_parent_build=parent)
 
-    def odcs_get_compose(self, compose_id):
-        """
-        Returns the information from the ODCS server about compose with id
-        `compose_id`. In DRY_RUN mode, returns fake compose information
-        without contacting the ODCS server.
-        """
-        if conf.dry_run:
-            compose = {}
-            compose['id'] = compose_id
-            compose['result_repofile'] = "http://localhost/%d.repo" % (
-                compose['id'])
-            compose['state'] = COMPOSE_STATES['done']
-            return compose
-
-        odcs = ODCS(conf.odcs_server_url, auth_mech=AuthMech.Kerberos,
-                    verify_ssl=conf.odcs_verify_ssl)
-        with krb_context():
-            return odcs.get_compose(compose_id)
-
     def get_repo_urls(self, db_event, build):
         """
         Returns list of URLs to ODCS repositories which should be used
diff --git a/freshmaker/handlers/errata/errata_advisory_rpms_signed.py b/freshmaker/handlers/errata/errata_advisory_rpms_signed.py
index 5ebce94..1843b56 100644
--- a/freshmaker/handlers/errata/errata_advisory_rpms_signed.py
+++ b/freshmaker/handlers/errata/errata_advisory_rpms_signed.py
@@ -27,7 +27,6 @@ import koji
 
 from freshmaker import conf, db, log
 from freshmaker.events import ErrataAdvisoryRPMsSignedEvent
-from freshmaker.events import ODCSComposeStateChangeEvent
 from freshmaker.handlers import BaseHandler, fail_event_on_handler_exception
 from freshmaker.kojiservice import koji_service
 from freshmaker.lightblue import LightBlue
@@ -35,12 +34,7 @@ from freshmaker.pulp import Pulp
 from freshmaker.errata import Errata
 from freshmaker.types import ArtifactType, ArtifactBuildState, EventState
 from freshmaker.models import Event
-from freshmaker.consumer import work_queue_put
-from freshmaker.utils import krb_context, retry, get_rebuilt_nvr
-
-from odcs.client.odcs import ODCS
-from odcs.client.odcs import AuthMech
-from odcs.common.types import COMPOSE_STATES
+from freshmaker.utils import get_rebuilt_nvr
 
 
 class ErrataAdvisoryRPMsSignedHandler(BaseHandler):
@@ -51,9 +45,6 @@ class ErrataAdvisoryRPMsSignedHandler(BaseHandler):
 
     name = 'ErrataAdvisoryRPMsSignedHandler'
 
-    # Used to generate incremental compose id in dry run mode.
-    _FAKE_COMPOSE_ID = 1
-
     def can_handle(self, event):
         return isinstance(event, ErrataAdvisoryRPMsSignedEvent)
 
@@ -145,35 +136,6 @@ class ErrataAdvisoryRPMsSignedHandler(BaseHandler):
 
         return []
 
-    def _fake_odcs_new_compose(self, compose_source, tag, packages=None):
-        """
-        Fake KojiSession.buildContainer method used dry run mode.
-
-        Logs the arguments and emits ErrataAdvisoryRPMsSignedHandler of
-        "done" state.
-
-        :rtype: dict
-        :return: Fake odcs.new_compose dict.
-        """
-        log.info("DRY RUN: Calling fake odcs.new_compose with args: %r",
-                 (compose_source, tag, packages))
-
-        # Generate the new_compose dict.
-        ErrataAdvisoryRPMsSignedHandler._FAKE_COMPOSE_ID += 1
-        new_compose = {}
-        new_compose['id'] = ErrataAdvisoryRPMsSignedHandler._FAKE_COMPOSE_ID
-        new_compose['result_repofile'] = "http://localhost/%d.repo" % (
-            new_compose['id'])
-        new_compose['state'] = COMPOSE_STATES['done']
-
-        # Generate and inject the ODCSComposeStateChangeEvent event.
-        event = ODCSComposeStateChangeEvent(
-            "fake_compose_msg", new_compose)
-        log.info("Injecting fake event: %r", event)
-        work_queue_put(event)
-
-        return new_compose
-
     def _prepare_yum_repo(self, db_event):
         """
         Prepare a yum repo for rebuild
@@ -208,21 +170,8 @@ class ErrataAdvisoryRPMsSignedHandler(BaseHandler):
                 % (builds, errata_id))
             return
 
-        log.info('Generate new compose for rebuild: '
-                 'source: %s, source type: %s, packages: %s',
-                 compose_source, 'tag', packages)
-
-        odcs = ODCS(conf.odcs_server_url, auth_mech=AuthMech.Kerberos,
-                    verify_ssl=conf.odcs_verify_ssl)
-        if not conf.dry_run:
-            with krb_context():
-                new_compose = odcs.new_compose(
-                    compose_source, 'tag', packages=packages,
-                    sigkeys=conf.odcs_sigkeys, flags=["no_deps"])
-        else:
-            new_compose = self._fake_odcs_new_compose(
-                compose_source, 'tag', packages=packages)
-
+        new_compose = self.odcs_new_compose("tag", compose_source,
+                                            packages=packages)
         compose_id = new_compose['id']
         yum_repourl = new_compose['result_repofile']
 
@@ -232,55 +181,22 @@ class ErrataAdvisoryRPMsSignedHandler(BaseHandler):
 
         return yum_repourl
 
-    def _prepare_pulp_repo(self, db_event, content_sets):
+    def _prepare_pulp_repo(self, build, content_sets):
         """
         Prepares .repo file containing the repositories matching
         the content_sets by creating new ODCS compose of PULP type.
 
         This currently blocks until the compose is done or failed.
 
-        :param db_event: models.Event instance associated with this build.
+        :param build: models.ArtifactBuild instance to build Pulp repo for.
         :param list content_sets: List of content sets.
         :rtype: dict
         :return: ODCS compose dictionary.
         """
-        log.info('Generating new PULP type compose for content_sets: %r',
-                 content_sets)
-
-        odcs = ODCS(conf.odcs_server_url, auth_mech=AuthMech.Kerberos,
-                    verify_ssl=conf.odcs_verify_ssl)
-        if not conf.dry_run:
-            with krb_context():
-                new_compose = odcs.new_compose(
-                    ' '.join(content_sets), 'pulp')
-
-                # Pulp composes in ODCS takes just few seconds, because ODCS
-                # only generates the .repo file after single query to Pulp.
-                # TODO: Freshmaker is currently not designed to handle
-                # multiple ODCS composes per rebuild Event and since these
-                # composes are done in no-time normally, it is OK here to
-                # block. It would still be nice to redesign that part of
-                # Freshmaker to do things "right".
-                # This is tracked here: https://pagure.io/freshmaker/issue/114
-                @retry(timeout=60, interval=2)
-                def wait_for_compose(compose_id):
-                    ret = odcs.get_compose(compose_id)
-                    if ret["state_name"] == "done":
-                        return True
-                    elif ret["state_name"] == "failed":
-                        return False
-                    log.info("Waiting for Pulp compose to finish: %r", ret)
-                    raise Exception("ODCS compose not finished.")
-
-                done = wait_for_compose(new_compose["id"])
-                if not done:
-                    db_event.builds_transition(
-                        ArtifactBuildState.FAILED.value, "Cannot generate "
-                        "ODCS PULP compose for content_sets %r"
-                        % (content_sets))
-        else:
-            new_compose = self._fake_odcs_new_compose(
-                content_sets, 'pulp')
+        new_compose = self.odcs_new_compose(
+            "pulp", ' '.join(content_sets), wait=True)
+        build.compose_id = new_compose["id"]
+        build.compose_done = True
 
         return new_compose
 
@@ -475,7 +391,7 @@ class ErrataAdvisoryRPMsSignedHandler(BaseHandler):
 
                 build.transition(state, state_reason)
 
-                compose = self._prepare_pulp_repo(build.event, image["content_sets"])
+                self._prepare_pulp_repo(build, image["content_sets"])
 
                 build_args = {}
                 build_args["repository"] = image["repository"]
@@ -483,7 +399,7 @@ class ErrataAdvisoryRPMsSignedHandler(BaseHandler):
                 build_args["parent"] = parent_nvr
                 build_args["target"] = image["target"]
                 build_args["branch"] = image["git_branch"]
-                build_args["odcs_pulp_compose_id"] = compose["id"]
+                build_args["content_sets"] = image["content_sets"]
                 build.build_args = json.dumps(build_args)
                 db.session.commit()
 
diff --git a/freshmaker/handlers/odcs/compose_state_change.py b/freshmaker/handlers/odcs/compose_state_change.py
index 4d3f35a..8fff83d 100644
--- a/freshmaker/handlers/odcs/compose_state_change.py
+++ b/freshmaker/handlers/odcs/compose_state_change.py
@@ -21,11 +21,12 @@
 #
 # Written by Chenxiong Qi <cqi@redhat.com>
 
-from freshmaker import db
-from freshmaker.models import Event
+from freshmaker import db, log
+from freshmaker.models import Event, ArtifactBuild
 from freshmaker.handlers import (
     ContainerBuildHandler, fail_event_on_handler_exception)
 from freshmaker.events import ODCSComposeStateChangeEvent
+import json
 
 from odcs.common.types import COMPOSE_STATES
 
@@ -46,4 +47,76 @@ class ComposeStateChangeHandler(ContainerBuildHandler):
             Event.compose_id == event.compose['id']).all()
         for db_event in errata_signed_events:
             self.set_context(db_event)
-            self._build_first_batch(db_event)
+
+            # For container images rebuilds, we always get two ODCS composes:
+            # - Pulp compose with latest shipped versions of packages. The
+            #   [repo] name of this compose in its .repo file is the same as
+            #   name of content_set from which compose was generated.
+            # - Koji tag compose with packages from Errata advisory with which
+            #   we rebuild the container image. The [repo] name of this compose
+            #   in its .repo file is "odcs-N" where N is ID of compose.
+            # We could just rebuild container image using these two ODCS
+            # composes directly, BUT some Dockerfiles do following:
+            #
+            # RUN yum-config-manager --disable \* &> /dev/null &&
+            #     yum-config-manager --enable rhel-7-server-rpms &&
+            #     yum-config-manager --enable rhel-7-server-optional-rpms &&
+            #     yum-config-manager --enable rhel-server-rhscl-7-rpms
+            #
+            # This means that they would disable our "odcs-N" repository and
+            # enable only the Pulp repositories. This means that RPMs from
+            # Errata advisory would never get in.
+            #
+            # To workaround that, we have to generate new compose which merges
+            # Pulp compose and odcs-N compose together, uses the latest NVRs
+            # of packages from both and keeps the [repo] name of the Pulp
+            # repository in the .repo file of resulting compose.
+            #
+            # That way, the "rhel-7-server-rpms" repository from example above
+            # contains packages from Pulp AND also so far unreleased
+            # packages from Errata advisory.
+            log.info("Main ODCS compose %d is done. Submitting ODCS tasks "
+                     "to merge Pulp composes with this compose.",
+                     int(event.compose["id"]))
+            for build in db_event.builds:
+                if not build.build_args or "content_sets" not in build.build_args:
+                    raise ValueError("content_sets not in build_args for "
+                                     "build %r" % build)
+
+                if not build.compose_id:
+                    raise ValueError("compose_id not set for "
+                                     "build %r" % build)
+
+                args = json.loads(build.build_args)
+
+                new_compose = self.odcs_new_compose(
+                    "odcs_compose",
+                    "%d %d" % (event.compose['id'], build.compose_id),
+                    result_repo_name=args["content_sets"][0])
+                build.compose_id = new_compose["id"]
+                build.compose_done = False
+                db.session.commit()
+
+        builds = db.session.query(ArtifactBuild).filter(
+            ArtifactBuild.compose_id == event.compose['id']).all()
+        if builds:
+            build = builds[0]
+            self.set_context(build.event)
+            build.compose_done = True
+            db.session.commit()
+
+            has_undone_composes = False
+            for b in build.event.builds:
+                if b.compose_id and not b.compose_done:
+                    has_undone_composes = True
+                    break
+
+            if has_undone_composes:
+                log.info("Merge ODCS compose %d is done, waiting for others "
+                         "composes associated with Event %r to finish",
+                         int(event.compose['id']), build.event)
+            else:
+                log.info("Merge ODCS compose %d is done as well as all others "
+                         "composes associated with Event %r to finish",
+                         int(event.compose['id']), build.event)
+                self._build_first_batch(build.event)
diff --git a/freshmaker/migrations/versions/f08127e4baad_.py b/freshmaker/migrations/versions/f08127e4baad_.py
new file mode 100644
index 0000000..48ce5d5
--- /dev/null
+++ b/freshmaker/migrations/versions/f08127e4baad_.py
@@ -0,0 +1,24 @@
+"""Add compose_done and compose_id to artifact_builds.
+
+Revision ID: f08127e4baad
+Revises: 90f8444d5ab7
+Create Date: 2017-12-04 13:46:46.031362
+
+"""
+
+# revision identifiers, used by Alembic.
+revision = 'f08127e4baad'
+down_revision = '90f8444d5ab7'
+
+from alembic import op
+import sqlalchemy as sa
+
+
+def upgrade():
+    op.add_column('artifact_builds', sa.Column('compose_done', sa.Boolean(), nullable=True))
+    op.add_column('artifact_builds', sa.Column('compose_id', sa.Integer(), nullable=True))
+
+
+def downgrade():
+    op.drop_column('artifact_builds', 'compose_id')
+    op.drop_column('artifact_builds', 'compose_done')
diff --git a/freshmaker/models.py b/freshmaker/models.py
index ff0b61f..34b33c0 100644
--- a/freshmaker/models.py
+++ b/freshmaker/models.py
@@ -320,6 +320,16 @@ class ArtifactBuild(FreshmakerBase):
     # Build args in json format.
     build_args = db.Column(db.String, nullable=True)
 
+    compose_id = db.Column(
+        db.Integer,
+        default=None,
+        doc='Used to track per-artifact-build ODCS compose.')
+
+    compose_done = db.Column(
+        db.Boolean,
+        default=False,
+        doc='True when compose defined by compose_id is done.')
+
     @classmethod
     def create(cls, session, event, name, type,
                build_id=None, dep_on=None, state=None,
diff --git a/tests/test_errata_advisory_state_changed.py b/tests/test_errata_advisory_state_changed.py
index 890b921..3dcd2df 100644
--- a/tests/test_errata_advisory_state_changed.py
+++ b/tests/test_errata_advisory_state_changed.py
@@ -274,9 +274,9 @@ class TestBatches(unittest.TestCase):
                 'parent': parent, "target": "t1", 'git_branch': 'mybranch',
                 "error": error, "content_sets": ["first-content-set"]}
 
-    @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.ODCS.new_compose')
-    @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.ODCS.get_compose')
-    @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.krb_context')
+    @patch('freshmaker.handlers.ODCS.new_compose')
+    @patch('freshmaker.handlers.ODCS.get_compose')
+    @patch('freshmaker.handlers.krb_context')
     def test_batches_records(self, krb_context, get_compose, new_compose):
         """
         Tests that batches are properly recorded in DB.
@@ -599,7 +599,7 @@ class TestPrepareYumRepo(unittest.TestCase):
         db.drop_all()
         db.session.commit()
 
-    @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.ODCS')
+    @patch('freshmaker.handlers.ODCS')
     @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.'
            'ErrataAdvisoryRPMsSignedHandler._get_packages_for_compose')
     @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.'
@@ -643,7 +643,7 @@ class TestPrepareYumRepo(unittest.TestCase):
             "http://localhost/composes/latest-odcs-3-1/compose/Temporary/odcs-3.repo",
             repo_url)
 
-    @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.ODCS')
+    @patch('freshmaker.handlers.ODCS')
     @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.'
            'ErrataAdvisoryRPMsSignedHandler._get_packages_for_compose')
     @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.'
@@ -674,7 +674,7 @@ class TestPrepareYumRepo(unittest.TestCase):
             self.assertEqual(build.state_reason, "Packages for errata "
                              "advisory 123 found in multiple different tags.")
 
-    @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.ODCS')
+    @patch('freshmaker.handlers.ODCS')
     @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.'
            'ErrataAdvisoryRPMsSignedHandler._get_packages_for_compose')
     @patch('freshmaker.handlers.errata.errata_advisory_rpms_signed.'
@@ -900,14 +900,14 @@ class TestRecordBatchesImages(unittest.TestCase):
                                               {self.mock_event.__class__: -1})
         self.event_types_patcher.start()
 
-        self.prepare_pulp_repo_patcher = patch(
+        self.odcs_new_compose_patcher = patch(
             'freshmaker.handlers.errata.'
-            'ErrataAdvisoryRPMsSignedHandler._prepare_pulp_repo',
+            'ErrataAdvisoryRPMsSignedHandler.odcs_new_compose',
             side_effect=[{'id': 1}, {'id': 2}])
-        self.mock_prepare_pulp_repo = self.prepare_pulp_repo_patcher.start()
+        self.odcs_new_compose = self.odcs_new_compose_patcher.start()
 
     def tearDown(self):
-        self.prepare_pulp_repo_patcher.stop()
+        self.odcs_new_compose_patcher.stop()
         self.event_types_patcher.stop()
 
         db.session.remove()
@@ -970,8 +970,8 @@ class TestRecordBatchesImages(unittest.TestCase):
         self.assertNotEqual(None, parent_image)
         self.assertEqual(ArtifactBuildState.PLANNED.value, parent_image.state)
 
-        build_args = json.loads(parent_image.build_args)
-        self.assertEqual(1, build_args['odcs_pulp_compose_id'])
+        self.assertEqual(parent_image.compose_id, 1)
+        self.assertEqual(parent_image.compose_done, True)
 
         # Check child image
         child_image = query.filter(
@@ -981,12 +981,12 @@ class TestRecordBatchesImages(unittest.TestCase):
         self.assertEqual(parent_image, child_image.dep_on)
         self.assertEqual(ArtifactBuildState.PLANNED.value, child_image.state)
 
-        build_args = json.loads(child_image.build_args)
-        self.assertEqual(2, build_args['odcs_pulp_compose_id'])
+        self.assertEqual(child_image.compose_id, 2)
+        self.assertEqual(child_image.compose_done, True)
 
-        self.mock_prepare_pulp_repo.assert_has_calls([
-            call(child_image.event, ["content-set-1"]),
-            call(child_image.event, ["content-set-1"])
+        self.odcs_new_compose.assert_has_calls([
+            call('pulp', 'content-set-1', wait=True),
+            call('pulp', 'content-set-1', wait=True)
         ])
 
     def test_mark_failed_state_if_image_has_error(self):
diff --git a/tests/test_odcs_compose_state_change.py b/tests/test_odcs_compose_state_change.py
index 63dde15..8898068 100644
--- a/tests/test_odcs_compose_state_change.py
+++ b/tests/test_odcs_compose_state_change.py
@@ -26,7 +26,7 @@ import unittest
 from mock import call, patch
 
 from freshmaker import db
-from freshmaker.models import Event
+from freshmaker.models import Event, ArtifactBuild
 from freshmaker.models import EVENT_TYPES
 from freshmaker.events import ErrataAdvisoryRPMsSignedEvent
 from freshmaker.events import KojiTaskStateChangeEvent
@@ -46,15 +46,26 @@ class TestComposeStateChangeHandler(unittest.TestCase):
         self.adv_signed_event1 = Event.get_or_create(
             db.session, 'msg-id-1', 'msg-id-1',
             EVENT_TYPES[ErrataAdvisoryRPMsSignedEvent])
+        self.adv_signed_event1_build1 = ArtifactBuild.create(
+            db.session, self.adv_signed_event1, "ed", "module", 1234)
+        self.adv_signed_event1_build2 = ArtifactBuild.create(
+            db.session, self.adv_signed_event1, "mksh", "module", 1235)
+
         self.adv_signed_event2 = Event.get_or_create(
             db.session, 'msg-id-2', 'msg-id-2',
             EVENT_TYPES[ErrataAdvisoryRPMsSignedEvent])
+        self.adv_signed_event2_build1 = ArtifactBuild.create(
+            db.session, self.adv_signed_event2, "ed", "module", 1234)
+        self.adv_signed_event2_build2 = ArtifactBuild.create(
+            db.session, self.adv_signed_event2, "mksh", "module", 1235)
+
         self.unrelated_event = Event.get_or_create(
             db.session, 'msg-id-3', 'msg-id-3',
             EVENT_TYPES[KojiTaskStateChangeEvent])
 
         self.adv_signed_event1.compose_id = 1
         self.adv_signed_event2.compose_id = 1
+
         db.session.commit()
 
     def tearDown(self):
@@ -72,17 +83,99 @@ class TestComposeStateChangeHandler(unittest.TestCase):
 
     @patch('freshmaker.handlers.ContainerBuildHandler._build_first_batch')
     @patch('freshmaker.handlers.ContainerBuildHandler.set_context')
-    def test_start_to_build(self, set_context, build_first_batch):
+    def test_start_to_build_no_content_sets(
+            self, set_context, build_first_batch):
+        event = ODCSComposeStateChangeEvent(
+            'msg-id', {'id': 1, 'state': 'done'}
+        )
+        handler = ComposeStateChangeHandler()
+        self.assertRaises(ValueError, handler.handle, event)
+
+    @patch('freshmaker.handlers.ContainerBuildHandler._build_first_batch')
+    @patch('freshmaker.handlers.ContainerBuildHandler.set_context')
+    def test_start_to_build_no_compose_id(
+            self, set_context, build_first_batch):
+        self.adv_signed_event1_build1.build_args = '{"content_sets": ["foo", "bar"]}'
+        self.adv_signed_event1_build2.build_args = '{"content_sets": ["foo", "bar"]}'
+        self.adv_signed_event2_build1.build_args = '{"content_sets": ["foo", "cs"]}'
+        self.adv_signed_event2_build2.build_args = '{"content_sets": ["foo", "cs"]}'
+        db.session.commit()
+
+        event = ODCSComposeStateChangeEvent(
+            'msg-id', {'id': 1, 'state': 'done'}
+        )
+        handler = ComposeStateChangeHandler()
+        self.assertRaises(ValueError, handler.handle, event)
+
+    @patch('freshmaker.handlers.ContainerBuildHandler.odcs_new_compose')
+    @patch('freshmaker.handlers.ContainerBuildHandler.set_context')
+    def test_start_to_build(
+            self, set_context, odcs_new_compose):
+        self.adv_signed_event1_build1.build_args = '{"content_sets": ["foo", "bar"]}'
+        self.adv_signed_event1_build2.build_args = '{"content_sets": ["foo", "bar"]}'
+        self.adv_signed_event2_build1.build_args = '{"content_sets": ["foo", "cs"]}'
+        self.adv_signed_event2_build2.build_args = '{"content_sets": ["foo", "cs"]}'
+        self.adv_signed_event1_build1.compose_id = 2
+        self.adv_signed_event1_build2.compose_id = 3
+        self.adv_signed_event2_build1.compose_id = 4
+        self.adv_signed_event2_build2.compose_id = 5
+        self.adv_signed_event1_build1.compose_done = True
+        self.adv_signed_event1_build2.compose_done = True
+        self.adv_signed_event2_build1.compose_done = True
+        self.adv_signed_event2_build2.compose_done = True
+        db.session.commit()
+
+        odcs_new_compose.return_value = {"id": 123}
+
         event = ODCSComposeStateChangeEvent(
             'msg-id', {'id': 1, 'state': 'done'}
         )
         handler = ComposeStateChangeHandler()
         handler.handle(event)
-        build_first_batch.assert_has_calls([
-            call(self.adv_signed_event1),
-            call(self.adv_signed_event2),
+
+        self.assertEqual(self.adv_signed_event1_build1.compose_id, 123)
+        self.assertEqual(self.adv_signed_event1_build2.compose_id, 123)
+        self.assertEqual(self.adv_signed_event2_build1.compose_id, 123)
+        self.assertEqual(self.adv_signed_event2_build2.compose_id, 123)
+        self.assertEqual(self.adv_signed_event1_build1.compose_done, False)
+        self.assertEqual(self.adv_signed_event1_build2.compose_done, False)
+        self.assertEqual(self.adv_signed_event2_build1.compose_done, False)
+        self.assertEqual(self.adv_signed_event2_build2.compose_done, False)
+
+        odcs_new_compose.assert_has_calls([
+            call('odcs_compose', '1 2', result_repo_name=u'foo'),
+            call('odcs_compose', '1 3', result_repo_name=u'foo'),
+            call('odcs_compose', '1 4', result_repo_name=u'foo'),
+            call('odcs_compose', '1 5', result_repo_name=u'foo'),
         ])
 
+    @patch('freshmaker.handlers.ContainerBuildHandler._build_first_batch')
+    @patch('freshmaker.handlers.ContainerBuildHandler.set_context')
+    def test_start_to_build_odcs_compose(
+            self, set_context, build_first_batch):
+        self.adv_signed_event1_build1.compose_id = 2
+        self.adv_signed_event1_build2.compose_id = 3
+        self.adv_signed_event2_build1.compose_id = 4
+        self.adv_signed_event2_build2.compose_id = 5
+
+        for i in [2, 3, 4, 5]:
+            event = ODCSComposeStateChangeEvent(
+                'msg-id', {'id': i, 'state': 'done'}
+            )
+            handler = ComposeStateChangeHandler()
+            handler.handle(event)
+
+            if i == 2:
+                build_first_batch.assert_not_called()
+            elif i in [3, 4]:
+                build_first_batch.assert_has_calls(
+                    [call(self.adv_signed_event1)])
+            else:
+                build_first_batch.assert_has_calls([
+                    call(self.adv_signed_event1),
+                    call(self.adv_signed_event2),
+                ])
+
         set_context.assert_has_calls([
             call(self.adv_signed_event1),
             call(self.adv_signed_event2)

