diff --git a/pungi/ostree.py b/pungi/ostree.py
deleted file mode 100644
index 16e7fc2d..00000000
--- a/pungi/ostree.py
+++ /dev/null
@@ -1,81 +0,0 @@
-# -*- coding: utf-8 -*-
-
-"""
-This module contains functions required by pungi-make-ostree.
-It is expected to be runnable in Koji runroot.
-"""
-
-import argparse
-import os
-from kobo import shortcuts
-import errno
-
-
-def ensure_dir(path):
-    try:
-        os.makedirs(path)
-    except OSError as err:
-        if err.errno != errno.EEXIST:
-            raise
-    return path
-
-
-def make_log_file(log_dir, filename):
-    """Return path to log file with given name, if log_dir is set."""
-    if not log_dir:
-        return None
-    ensure_dir(log_dir)
-    return os.path.join(log_dir, '%s.log' % filename)
-
-
-def init_ostree_repo(repo, log_dir=None):
-    """If the ostree repo does not exist, initialize it."""
-    log_file = make_log_file(log_dir, 'init-ostree-repo')
-    if not os.path.isdir(repo) or not os.listdir(repo):
-        ensure_dir(repo)
-        shortcuts.run(['ostree', 'init', '--repo=%s' % repo, '--mode=archive-z2'],
-                      show_cmd=True, stdout=True, logfile=log_file)
-
-
-def make_ostree_repo(repo, config, version=None, log_dir=None):
-    log_file = make_log_file(log_dir, 'create-ostree-repo')
-    cmd = ['rpm-ostree', 'compose', 'tree', '--repo=%s' % repo,
-           '--write-commitid-to=%s' % make_log_file(log_dir, 'commitid')]
-    if version:
-        # Add versioning metadata
-        cmd.append('--add-metadata-string=version=%s' % version)
-    cmd.append(config)
-
-    shortcuts.run(cmd, show_cmd=True, stdout=True, logfile=log_file)
-
-
-def update_ostree_summary(repo, log_dir=None):
-    log_file = make_log_file(log_dir, 'ostree-summary')
-    shortcuts.run(['ostree', 'summary', '-u', '--repo=%s' % repo],
-                  show_cmd=True, stdout=True, logfile=log_file)
-
-
-def run(opts):
-    init_ostree_repo(opts.ostree_repo, log_dir=opts.log_dir)
-    make_ostree_repo(opts.ostree_repo, opts.treefile, version=opts.version, log_dir=opts.log_dir)
-    if opts.update_summary:
-        update_ostree_summary(opts.ostree_repo, log_dir=opts.log_dir)
-
-
-def main(args=None):
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--log-dir',
-                        help='where to log output')
-
-    parser.add_argument('ostree_repo', metavar='OSTREE_REPO',
-                        help='where to put the ostree repo')
-    parser.add_argument('--treefile', required=True,
-                        help='treefile for rpm-ostree')
-    parser.add_argument('--version',
-                        help='version string to be added as versioning metadata')
-    parser.add_argument('--update-summary', action='store_true',
-                        help='update summary metadata')
-
-    opts = parser.parse_args(args)
-
-    run(opts)
diff --git a/pungi/ostree/__init__.py b/pungi/ostree/__init__.py
new file mode 100644
index 00000000..fe8c53ea
--- /dev/null
+++ b/pungi/ostree/__init__.py
@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+
+
+# 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
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Library General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, see <https://gnu.org/licenses/>.
+
+
+import argparse
+
+from .tree import Tree
+
+
+def main(args=None):
+    parser = argparse.ArgumentParser()
+    subparser = parser.add_subparsers(help="Sub commands")
+
+    treep = subparser.add_parser("tree", help="Compose OSTree repository")
+    treep.set_defaults(_class=Tree, func='run')
+    treep.add_argument('--repo', metavar='PATH', required=True,
+                       help='where to put the OSTree repo (required)')
+    treep.add_argument('--treefile', metavar="FILE", required=True,
+                       help='treefile for rpm-ostree (required)')
+    treep.add_argument('--log-dir', metavar="DIR",
+                       help='where to log output')
+    treep.add_argument('--extra-config', metavar="FILE",
+                       help='JSON file contains extra configurations')
+    treep.add_argument('--version', metavar="VERSION",
+                       help='version string to be added as versioning metadata')
+    treep.add_argument('--update-summary', action='store_true',
+                       help='update summary metadata')
+
+    args = parser.parse_args(args)
+    _class = args._class()
+    _class.set_args(args)
+    func = getattr(_class, args.func)
+    func()
diff --git a/pungi/ostree/base.py b/pungi/ostree/base.py
new file mode 100644
index 00000000..4932f597
--- /dev/null
+++ b/pungi/ostree/base.py
@@ -0,0 +1,19 @@
+# -*- coding: utf-8 -*-
+
+
+# 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
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Library General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, see <https://gnu.org/licenses/>.
+
+
+class OSTree(object):
+    def set_args(self, args):
+        self.args = args
diff --git a/pungi/ostree/tree.py b/pungi/ostree/tree.py
new file mode 100644
index 00000000..1b616b2c
--- /dev/null
+++ b/pungi/ostree/tree.py
@@ -0,0 +1,101 @@
+# -*- coding: utf-8 -*-
+
+
+# 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
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Library General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, see <https://gnu.org/licenses/>.
+
+
+import os
+import json
+from kobo import shortcuts
+
+from pungi.util import makedirs
+from .base import OSTree
+from .utils import (make_log_file, tweak_treeconf,
+                    get_ref_from_treefile, get_commitid_from_commitid_file)
+
+
+class Tree(OSTree):
+    def _init_repo(self):
+        """If the ostree repo does not exist, initialize it."""
+        log_file = make_log_file(self.logdir, 'init-ostree-repo')
+        if not os.path.isdir(self.repo) or not os.listdir(self.repo):
+            makedirs(self.repo)
+            shortcuts.run(['ostree', 'init', '--repo=%s' % self.repo, '--mode=archive-z2'],
+                          show_cmd=True, stdout=True, logfile=log_file)
+
+    def _make_tree(self):
+        """Compose OSTree tree"""
+        log_file = make_log_file(self.logdir, 'create-ostree-repo')
+        cmd = ['rpm-ostree', 'compose', 'tree', '--repo=%s' % self.repo,
+               '--write-commitid-to=%s' % self.commitid_file]
+        if self.version:
+            # Add versioning metadata
+            cmd.append('--add-metadata-string=version=%s' % self.version)
+        cmd.append(self.treefile)
+
+        shortcuts.run(cmd, show_cmd=True, stdout=True, logfile=log_file)
+
+    def _update_summary(self):
+        """Update summary metadata"""
+        log_file = make_log_file(self.logdir, 'ostree-summary')
+        shortcuts.run(['ostree', 'summary', '-u', '--repo=%s' % self.repo],
+                      show_cmd=True, stdout=True, logfile=log_file)
+
+    def _update_ref(self):
+        """
+        Update the ref.
+
+        '--write-commitid-to' is specified when compose the tree, so we need
+        to update the ref by ourselves. ref is retrieved from treefile and
+        commitid is retrieved from the committid file.
+        """
+        tag_ref = True
+        if self.extra_config:
+            tag_ref = self.extra_config.get('tag_ref', True)
+        if not tag_ref:
+            return
+        ref = get_ref_from_treefile(self.treefile)
+        commitid = get_commitid_from_commitid_file(self.commitid_file)
+        if ref and commitid:
+            # Let's write the tag out ourselves
+            heads_dir = os.path.join(self.repo, 'refs', 'heads')
+            if not os.path.exists(heads_dir):
+                raise RuntimeError('Refs/heads did not exist in ostree repo')
+
+            ref_path = os.path.join(heads_dir, ref)
+            makedirs(os.path.dirname(ref_path))
+            with open(ref_path, 'w') as f:
+                f.write(commitid + '\n')
+
+    def run(self):
+        self.repo = self.args.repo
+        self.treefile = self.args.treefile
+        self.version = self.args.version
+        self.logdir = self.args.log_dir
+        self.update_summary = self.args.update_summary
+        self.extra_config = self.args.extra_config
+        if self.extra_config:
+            self.extra_config = json.load(open(self.extra_config, 'r'))
+            source_repo_from = self.extra_config.get('source_repo_from', None)
+            extra_source_repos = self.extra_config.get('extra_source_repos', None)
+            keep_original_sources = self.extra_config.get('keep_original_sources', False)
+            repos = extra_source_repos + [{'name': 'source_repo_from', 'baseurl': source_repo_from}]
+            tweak_treeconf(self.treefile, source_repos=repos, keep_original_sources=keep_original_sources)
+
+        self.commitid_file = make_log_file(self.logdir, 'commitid')
+
+        self._init_repo()
+        self._make_tree()
+        self._update_ref()
+        if self.update_summary:
+            self._update_summary()
diff --git a/pungi/ostree/utils.py b/pungi/ostree/utils.py
new file mode 100644
index 00000000..fb5e3300
--- /dev/null
+++ b/pungi/ostree/utils.py
@@ -0,0 +1,93 @@
+# -*- coding: utf-8 -*-
+
+
+# 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
+# the Free Software Foundation; version 2 of the License.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Library General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program; if not, see <https://gnu.org/licenses/>.
+
+
+import datetime
+import json
+import os
+
+from pungi.util import makedirs
+
+
+def make_log_file(log_dir, filename):
+    """Return path to log file with given name, if log_dir is set."""
+    if not log_dir:
+        return None
+    makedirs(log_dir)
+    return os.path.join(log_dir, '%s.log' % filename)
+
+
+def get_ref_from_treefile(treefile):
+    """Return ref name by parsing the tree config file"""
+    ref = None
+    if os.path.isfile(treefile):
+        with open(treefile, 'r') as f:
+            try:
+                parsed = json.loads(f.read())
+                ref = parsed['ref']
+            except Exception:
+                pass
+    return ref
+
+
+def get_commitid_from_commitid_file(commitid_file):
+    """Return commit id which is read from the commitid file"""
+    commitid = None
+    if os.path.isfile(commitid_file):
+        with open(commitid_file, 'r') as f:
+            commitid = f.read().replace('\n', '')
+    return commitid
+
+
+def tweak_treeconf(treeconf, source_repos=None, keep_original_sources=False):
+    """
+    Update tree config file by adding new repos, and remove existing repos
+    from the tree config file if 'keep_original_sources' is not enabled.
+    """
+    # add this timestamp to repo name to get unique repo filename and repo name
+    # should be safe enough
+    time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
+
+    treeconf_dir = os.path.dirname(treeconf)
+    with open(treeconf, 'r') as f:
+        treeconf_content = json.load(f)
+
+    # backup the old tree config
+    os.rename(treeconf, '%s.%s.bak' % (treeconf, time))
+
+    repos = []
+    if source_repos:
+        for repo in source_repos:
+            name = "%s-%s" % (repo['name'], time)
+            with open("%s/%s.repo" % (treeconf_dir, name), 'w') as f:
+                f.write("[%s]\n" % name)
+                f.write("name=%s\n" % name)
+                f.write("baseurl=%s\n" % repo['baseurl'])
+                exclude = repo.get('exclude', None)
+                if exclude:
+                    f.write("exclude=%s\n" % exclude)
+                gpgcheck = '1' if repo.get('gpgcheck', False) else '0'
+                f.write("gpgcheck=%s\n" % gpgcheck)
+            repos.append(name)
+
+    original_repos = treeconf_content.get('repos', [])
+    if keep_original_sources:
+        treeconf_content['repos'] = original_repos + repos
+    else:
+        treeconf_content['repos'] = repos
+
+    # update tree config to add new repos
+    with open(treeconf, 'w') as f:
+        json.dump(treeconf_content, f, indent=4)
diff --git a/pungi/phases/ostree.py b/pungi/phases/ostree.py
index bd45cd2f..ebe03d0d 100644
--- a/pungi/phases/ostree.py
+++ b/pungi/phases/ostree.py
@@ -1,12 +1,13 @@
 # -*- coding: utf-8 -*-
 
-import datetime
+import copy
 import json
 import os
 from kobo.threads import ThreadPool, WorkerThread
 
 from .base import ConfigGuardedPhase
 from .. import util
+from ..ostree.utils import get_ref_from_treefile, get_commitid_from_commitid_file
 from ..paths import translate_path
 from ..wrappers import kojiwrapper, scm
 
@@ -53,7 +54,6 @@ class OSTreeThread(WorkerThread):
 
         self._clone_repo(repodir, config['config_url'], config.get('config_branch', 'master'))
 
-        treeconf = os.path.join(repodir, config['treefile'])
         source_repos = [{'name': '%s-%s' % (compose.compose_id, config['source_repo_from']),
                          'baseurl': source_repo}]
 
@@ -72,30 +72,41 @@ class OSTreeThread(WorkerThread):
 
             source_repos = source_repos + extra_source_repos
 
-        keep_original_sources = config.get('keep_original_sources', False)
-        self._tweak_treeconf(treeconf, source_repos=source_repos,
-                             keep_original_sources=keep_original_sources)
+        # copy the original config and update before save to a json file
+        new_config = copy.copy(config)
+
+        # repos in configuration can have repo url set to variant UID,
+        # update it to have the actual url that we just translated.
+        new_config.update({'source_repo_from': source_repo})
+        if extra_source_repos:
+            new_config.update({'extra_source_repos': extra_source_repos})
+
+        # remove unnecessary (for 'pungi-make-ostree tree' script ) elements
+        # from config, it doesn't hurt to have them, however remove them can
+        # reduce confusion
+        for k in ['ostree_repo', 'treefile', 'config_url', 'config_branch',
+                  'failable', 'version', 'update_summary']:
+            new_config.pop(k, None)
+
+        extra_config_file = None
+        if new_config:
+            # write a json file to save the configuration, so 'pungi-make-ostree tree'
+            # can take use of it
+            extra_config_file = os.path.join(workdir, 'extra_config.json')
+            with open(extra_config_file, 'w') as f:
+                json.dump(new_config, f, indent=4)
 
         # Ensure target directory exists, otherwise Koji task will fail to
         # mount it.
         util.makedirs(config['ostree_repo'])
 
-        self._run_ostree_cmd(compose, variant, arch, config, repodir)
-        ref, commitid = self._get_commit_info(config, repodir)
-        if config.get('tag_ref', True) and ref and commitid:
-            # Let's write the tag out ourselves
-            heads_dir = os.path.join(config['ostree_repo'], 'refs', 'heads')
-            if not os.path.exists(heads_dir):
-                raise RuntimeError('Refs/heads did not exist in ostree repo')
-
-            ref_path = os.path.join(heads_dir, ref)
-            if not os.path.exists(os.path.dirname(ref_path)):
-                os.makedirs(os.path.dirname(ref_path))
-
-            with open(ref_path, 'w') as f:
-                f.write(commitid + '\n')
+        self._run_ostree_cmd(compose, variant, arch, config, repodir,
+                             extra_config_file=extra_config_file)
 
         if compose.notifier:
+            ref = get_ref_from_treefile(os.path.join(repodir, config['treefile']))
+            # 'pungi-make-ostree tree' writes commitid to commitid.log in logdir
+            commitid = get_commitid_from_commitid_file(os.path.join(self.logdir, 'commitid.log'))
             compose.notifier.send('ostree',
                                   variant=variant.uid,
                                   arch=arch,
@@ -104,26 +115,12 @@ class OSTreeThread(WorkerThread):
 
         self.pool.log_info('[DONE ] %s' % msg)
 
-    def _get_commit_info(self, config, config_repo):
-        ref = None
-        commitid = None
-        with open(os.path.join(config_repo, config['treefile']), 'r') as f:
-            try:
-                parsed = json.loads(f.read())
-                ref = parsed['ref']
-            except ValueError:
-                return None, None
-        if os.path.exists(os.path.join(self.logdir, 'commitid')):
-            with open(os.path.join(self.logdir, 'commitid'), 'r') as f:
-                commitid = f.read().replace('\n', '')
-        else:
-            return None, None
-        return ref, commitid
-
-    def _run_ostree_cmd(self, compose, variant, arch, config, config_repo):
+    def _run_ostree_cmd(self, compose, variant, arch, config, config_repo, extra_config_file=None):
         cmd = [
             'pungi-make-ostree',
-            '--log-dir=%s' % os.path.join(self.logdir),
+            'tree',
+            '--repo=%s' % config['ostree_repo'],
+            '--log-dir=%s' % self.logdir,
             '--treefile=%s' % os.path.join(config_repo, config['treefile']),
         ]
 
@@ -131,12 +128,12 @@ class OSTreeThread(WorkerThread):
         if version:
             cmd.append('--version=%s' % version)
 
+        if extra_config_file:
+            cmd.append('--extra-config=%s' % extra_config_file)
+
         if config.get('update_summary', False):
             cmd.append('--update-summary')
 
-        # positional argument: ostree_repo
-        cmd.append(config['ostree_repo'])
-
         runroot_channel = compose.conf.get("runroot_channel")
         runroot_tag = compose.conf["runroot_tag"]
 
@@ -157,43 +154,3 @@ class OSTreeThread(WorkerThread):
     def _clone_repo(self, repodir, url, branch):
         scm.get_dir_from_scm({'scm': 'git', 'repo': url, 'branch': branch, 'dir': '.'},
                              repodir, logger=self.pool._logger)
-
-    def _tweak_treeconf(self, treeconf, source_repos, keep_original_sources=False):
-        """
-        Update tree config file by adding new repos and remove existing repos
-        from the tree config file if 'keep_original_sources' is not enabled.
-        """
-        # add this timestamp to repo name to get unique repo filename and repo name
-        # should be safe enough
-        time = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
-
-        treeconf_dir = os.path.dirname(treeconf)
-        with open(treeconf, 'r') as f:
-            treeconf_content = json.load(f)
-
-        # backup the old tree config
-        os.rename(treeconf, '%s.%s.bak' % (treeconf, time))
-
-        repos = []
-        for repo in source_repos:
-            name = "%s-%s" % (repo['name'], time)
-            with open("%s/%s.repo" % (treeconf_dir, name), 'w') as f:
-                f.write("[%s]\n" % name)
-                f.write("name=%s\n" % name)
-                f.write("baseurl=%s\n" % repo['baseurl'])
-                exclude = repo.get('exclude', None)
-                if exclude:
-                    f.write("exclude=%s\n" % exclude)
-                gpgcheck = '1' if repo.get('gpgcheck', False) else '0'
-                f.write("gpgcheck=%s\n" % gpgcheck)
-            repos.append(name)
-
-        original_repos = treeconf_content.get('repos', [])
-        if keep_original_sources:
-            treeconf_content['repos'] = original_repos + repos
-        else:
-            treeconf_content['repos'] = repos
-
-        # update tree config to add new repos
-        with open(treeconf, 'w') as f:
-            json.dump(treeconf_content, f, indent=4)
diff --git a/tests/test_ostree_phase.py b/tests/test_ostree_phase.py
index bf01f6c1..0861d9e5 100644
--- a/tests/test_ostree_phase.py
+++ b/tests/test_ostree_phase.py
@@ -110,10 +110,12 @@ class OSTreeThreadTest(helpers.PungiTestCase):
         self.assertEqual(koji.get_runroot_cmd.call_args_list,
                          [mock.call('rrt', 'x86_64',
                                     ['pungi-make-ostree',
+                                     'tree',
+                                     '--repo=%s' % self.repo,
                                      '--log-dir=%s/logs/x86_64/Everything/ostree-1' % self.topdir,
                                      '--treefile=%s/fedora-atomic-docker-host.json' % (
                                          self.topdir + '/work/ostree-1/config_repo'),
-                                     self.repo],
+                                     '--extra-config=%s/extra_config.json' % (self.topdir + '/work/ostree-1')],
                                     channel=None, mounts=[self.topdir, self.repo],
                                     packages=['pungi', 'ostree', 'rpm-ostree'],
                                     task_id=True, use_shell=True, new_chroot=True)])
@@ -121,17 +123,7 @@ class OSTreeThreadTest(helpers.PungiTestCase):
                          [mock.call(koji.get_runroot_cmd.return_value,
                                     log_file=self.topdir + '/logs/x86_64/Everything/ostree-1/runroot.log')])
 
-        repo_files = []
-        for fp in os.listdir(os.path.join(self.topdir, 'work/ostree-1/config_repo')):
-            if fp.endswith('.repo'):
-                repo_files.append(fp)
-
-                if fp not in ['fedora-rawhide.repo', 'fedora-24.repo', 'fedora-23.repo']:
-                    with open(os.path.join(self.topdir, 'work/ostree-1/config_repo', fp)) as f:
-                        self.assertIn('baseurl=http://example.com/Everything/$basearch/os', f.read())
-        # test a new repo file created
-        self.assertEqual(len(repo_files), 4)
-
+        self.assertTrue(os.path.isfile(os.path.join(self.topdir, 'work/ostree-1/extra_config.json')))
         self.assertTrue(os.path.isdir(self.repo))
 
     @mock.patch('pungi.wrappers.scm.get_dir_from_scm')
@@ -181,7 +173,7 @@ class OSTreeThreadTest(helpers.PungiTestCase):
         koji = KojiWrapper.return_value
         koji.run_runroot_cmd.side_effect = self._mock_runroot(
             0,
-            {'commitid': 'fca3465861a',
+            {'commitid.log': 'fca3465861a',
              'create-ostree-repo.log':
                 ['Doing work', 'fedora-atomic/25/x86_64 -> fca3465861a']})
         t = ostree.OSTreeThread(self.pool)
@@ -214,7 +206,7 @@ class OSTreeThreadTest(helpers.PungiTestCase):
                          [mock.call('ostree',
                                     variant='Everything',
                                     arch='x86_64',
-                                    ref=None,
+                                    ref='fedora-atomic/25/x86_64',
                                     commitid=None)])
 
     @mock.patch('pungi.wrappers.scm.get_dir_from_scm')
@@ -254,10 +246,13 @@ class OSTreeThreadTest(helpers.PungiTestCase):
         self.assertEqual(koji.get_runroot_cmd.call_args_list,
                          [mock.call('rrt', 'x86_64',
                                     ['pungi-make-ostree',
+                                     'tree',
+                                     '--repo=%s' % self.repo,
                                      '--log-dir=%s/logs/x86_64/Everything/ostree-1' % self.topdir,
                                      '--treefile=%s/fedora-atomic-docker-host.json' % (
                                          self.topdir + '/work/ostree-1/config_repo'),
-                                     '--update-summary', self.repo],
+                                     '--extra-config=%s/work/ostree-1/extra_config.json' % self.topdir,
+                                     '--update-summary'],
                                     channel=None, mounts=[self.topdir, self.repo],
                                     packages=['pungi', 'ostree', 'rpm-ostree'],
                                     task_id=True, use_shell=True, new_chroot=True)])
@@ -286,10 +281,13 @@ class OSTreeThreadTest(helpers.PungiTestCase):
         self.assertEqual(koji.get_runroot_cmd.call_args_list,
                          [mock.call('rrt', 'x86_64',
                                     ['pungi-make-ostree',
+                                     'tree',
+                                     '--repo=%s' % self.repo,
                                      '--log-dir=%s/logs/x86_64/Everything/ostree-1' % self.topdir,
                                      '--treefile=%s/fedora-atomic-docker-host.json' % (
                                          self.topdir + '/work/ostree-1/config_repo'),
-                                     '--version=24', self.repo],
+                                     '--version=24',
+                                     '--extra-config=%s/work/ostree-1/extra_config.json' % self.topdir],
                                     channel=None, mounts=[self.topdir, self.repo],
                                     packages=['pungi', 'ostree', 'rpm-ostree'],
                                     task_id=True, use_shell=True, new_chroot=True)])
@@ -299,7 +297,7 @@ class OSTreeThreadTest(helpers.PungiTestCase):
 
     @mock.patch('pungi.wrappers.scm.get_dir_from_scm')
     @mock.patch('pungi.wrappers.kojiwrapper.KojiWrapper')
-    def test_run_with_extra_source_repos(self, KojiWrapper, get_dir_from_scm):
+    def test_write_extra_config_file(self, KojiWrapper, get_dir_from_scm):
         get_dir_from_scm.side_effect = self._dummy_config_repo
 
         koji = KojiWrapper.return_value
@@ -319,55 +317,6 @@ class OSTreeThreadTest(helpers.PungiTestCase):
                     'exclude': 'systemd-container'
                 }
             ],
-            'config_url': 'https://git.fedorahosted.org/git/fedora-atomic.git',
-            'config_branch': 'f24',
-            'treefile': 'fedora-atomic-docker-host.json',
-            'ostree_repo': self.repo
-        }
-
-        t = ostree.OSTreeThread(self.pool)
-
-        t.process((self.compose, self.compose.variants['Everything'], 'x86_64', cfg), 1)
-
-        repo_files = []
-        for fp in os.listdir(os.path.join(self.topdir, 'work/ostree-1/config_repo')):
-            if fp.endswith('.repo'):
-                repo_files.append(fp)
-
-                if fp not in ['fedora-rawhide.repo', 'fedora-24.repo', 'fedora-23.repo']:
-                    if fp.startswith('repo_a'):
-                        with open(os.path.join(self.topdir, 'work/ostree-1/config_repo', fp)) as f:
-                            # ignore timestamp in repo name while checking
-                            content = f.read()
-                            self.assertIn('[repo_a', content)
-                            self.assertIn('name=repo_a', content)
-                            self.assertIn('baseurl=http://url/to/repo/a', content)
-                            self.assertIn('exclude=systemd-container', content)
-                            self.assertIn('gpgcheck=0', content)
-                    elif fp.startswith('Server'):
-                        with open(os.path.join(self.topdir, 'work/ostree-1/config_repo', fp)) as f:
-                            content = f.read()
-                            self.assertIn('[Server', content)
-                            self.assertIn('baseurl=http://example.com/Server/$basearch/os', content)
-                            self.assertIn('exclude=systemd-container', content)
-                            self.assertIn('gpgcheck=0', content)
-                    else:
-                        # this is the Everything repo (source_repo_from)
-                        with open(os.path.join(self.topdir, 'work/ostree-1/config_repo', fp)) as f:
-                            self.assertIn('baseurl=http://example.com/Everything/$basearch/os', f.read())
-        # test new repos files created
-        self.assertEqual(len(repo_files), 3 + 1 + len(cfg['extra_source_repos']))
-
-    @mock.patch('pungi.wrappers.scm.get_dir_from_scm')
-    @mock.patch('pungi.wrappers.kojiwrapper.KojiWrapper')
-    def test_run_with_keep_original_source_repos(self, KojiWrapper, get_dir_from_scm):
-        get_dir_from_scm.side_effect = self._dummy_config_repo
-
-        koji = KojiWrapper.return_value
-        koji.run_runroot_cmd.side_effect = self._mock_runroot(0)
-
-        cfg = {
-            'source_repo_from': 'Everything',
             'keep_original_sources': True,
             'config_url': 'https://git.fedorahosted.org/git/fedora-atomic.git',
             'config_branch': 'f24',
@@ -379,12 +328,13 @@ class OSTreeThreadTest(helpers.PungiTestCase):
 
         t.process((self.compose, self.compose.variants['Everything'], 'x86_64', cfg), 1)
 
-        treeconf_content = json.load(open(os.path.join(self.topdir,
-                                                       'work/ostree-1/config_repo',
-                                                       cfg['treefile'])))
-
-        # added 1 repo (Everything), have 4 (3 + 1) repos now
-        self.assertEqual(len(treeconf_content['repos']), 4)
+        extra_config_file = os.path.join(self.topdir, 'work/ostree-1/extra_config.json')
+        self.assertTrue(os.path.isfile(extra_config_file))
+        extra_config = json.load(open(extra_config_file, 'r'))
+        self.assertTrue(extra_config.get('keep_original_sources', False))
+        self.assertEqual(extra_config.get('source_repo_from', None), 'http://example.com/Everything/$basearch/os')
+        self.assertEqual(len(extra_config.get('extra_source_repos', [])), len(cfg['extra_source_repos']))
+        self.assertEqual(extra_config.get('extra_source_repos').pop()['baseurl'], 'http://example.com/Server/$basearch/os')
 
 if __name__ == '__main__':
     unittest.main()
diff --git a/tests/test_ostree_script.py b/tests/test_ostree_script.py
index 89b18268..735a8fb7 100644
--- a/tests/test_ostree_script.py
+++ b/tests/test_ostree_script.py
@@ -6,7 +6,9 @@ import unittest
 import mock
 
 import os
+import json
 import sys
+import datetime
 
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))
 sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'bin'))
@@ -15,16 +17,28 @@ from tests import helpers
 from pungi import ostree
 
 
-class OstreeScriptTest(helpers.PungiTestCase):
+class OstreeTreeScriptTest(helpers.PungiTestCase):
+
+    def _make_dummy_config_dir(self, path):
+        helpers.touch(os.path.join(path, 'fedora-atomic-docker-host.json'),
+                      json.dumps({'ref': 'fedora-atomic/25/x86_64',
+                                  'repos': ['fedora-rawhide', 'fedora-24', 'fedora-23']}))
+        helpers.touch(os.path.join(path, 'fedora-rawhide.repo'),
+                      '[fedora-rawhide]\nmirrorlist=mirror-mirror-on-the-wall')
+        helpers.touch(os.path.join(path, 'fedora-24.repo'),
+                      '[fedora-24]\nmetalink=who-is-the-fairest-of-them-all')
+        helpers.touch(os.path.join(path, 'fedora-23.repo'),
+                      '[fedora-23]\nbaseurl=why-not-zoidberg?')
 
     @mock.patch('kobo.shortcuts.run')
     def test_full_run(self, run):
         repo = os.path.join(self.topdir, 'atomic')
 
         ostree.main([
+            'tree',
+            '--repo=%s' % repo,
             '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
             '--treefile=%s/fedora-atomic-docker-host.json' % self.topdir,
-            repo,
         ])
 
         self.maxDiff = None
@@ -44,9 +58,10 @@ class OstreeScriptTest(helpers.PungiTestCase):
         os.mkdir(repo)
 
         ostree.main([
+            'tree',
+            '--repo=%s' % repo,
             '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
             '--treefile=%s/fedora-atomic-docker-host.json' % self.topdir,
-            repo,
         ])
 
         self.maxDiff = None
@@ -66,9 +81,10 @@ class OstreeScriptTest(helpers.PungiTestCase):
         helpers.touch(os.path.join(repo, 'initialized'))
 
         ostree.main([
+            'tree',
+            '--repo=%s' % repo,
             '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
             '--treefile=%s/fedora-atomic-docker-host.json' % self.topdir,
-            repo,
         ])
 
         self.maxDiff = None
@@ -84,10 +100,11 @@ class OstreeScriptTest(helpers.PungiTestCase):
         repo = os.path.join(self.topdir, 'atomic')
 
         ostree.main([
+            'tree',
+            '--repo=%s' % repo,
             '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
             '--treefile=%s/fedora-atomic-docker-host.json' % self.topdir,
             '--update-summary',
-            repo,
         ])
 
         self.maxDiff = None
@@ -107,10 +124,11 @@ class OstreeScriptTest(helpers.PungiTestCase):
         repo = os.path.join(self.topdir, 'atomic')
 
         ostree.main([
+            'tree',
+            '--repo=%s' % repo,
             '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
             '--treefile=%s/fedora-atomic-docker-host.json' % self.topdir,
             '--version=24',
-            repo,
         ])
 
         self.maxDiff = None
@@ -124,5 +142,129 @@ class OstreeScriptTest(helpers.PungiTestCase):
                         self.topdir + '/fedora-atomic-docker-host.json'],
                        logfile=self.topdir + '/logs/Atomic/create-ostree-repo.log', show_cmd=True, stdout=True)])
 
+    @mock.patch('pungi.ostree.utils.datetime')
+    @mock.patch('kobo.shortcuts.run')
+    def test_extra_config_with_extra_repos(self, run, time):
+        time.datetime.now.return_value = datetime.datetime(2016, 1, 1, 1, 1)
+        timestamp = time.datetime.now().strftime("%Y%m%d%H%M%S")
+
+        configdir = os.path.join(self.topdir, 'config')
+        self._make_dummy_config_dir(configdir)
+        treefile = os.path.join(configdir, 'fedora-atomic-docker-host.json')
+
+        repo = os.path.join(self.topdir, 'atomic')
+
+        extra_config_file = os.path.join(self.topdir, 'extra_config.json')
+        extra_config = {
+            "source_repo_from": "http://www.example.com/Server.repo",
+            "extra_source_repos": [
+                {
+                    "name": "optional",
+                    "baseurl": "http://example.com/repo/x86_64/optional",
+                    "exclude": "systemd-container",
+                    "gpgcheck": False
+                },
+                {
+                    "name": "extra",
+                    "baseurl": "http://example.com/repo/x86_64/extra",
+                }
+            ]
+        }
+        helpers.touch(extra_config_file, json.dumps(extra_config))
+
+        ostree.main([
+            'tree',
+            '--repo=%s' % repo,
+            '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
+            '--treefile=%s' % treefile,
+            '--extra-config=%s' % extra_config_file,
+        ])
+
+        source_repo_from_name = "source_repo_from-%s" % timestamp
+        source_repo_from_repo = os.path.join(configdir, "%s.repo" % source_repo_from_name)
+        self.assertTrue(os.path.isfile(source_repo_from_repo))
+        with open(source_repo_from_repo, 'r') as f:
+            content = f.read()
+            self.assertIn("[%s]" % source_repo_from_name, content)
+            self.assertIn("name=%s" % source_repo_from_name, content)
+            self.assertIn("baseurl=http://www.example.com/Server.repo", content)
+            self.assertIn("gpgcheck=0", content)
+
+        optional_repo_name = "optional-%s" % timestamp
+        optional_repo = os.path.join(configdir, "%s.repo" % optional_repo_name)
+        self.assertTrue(os.path.isfile(optional_repo))
+        with open(optional_repo, 'r') as f:
+            content = f.read()
+            self.assertIn("[%s]" % optional_repo_name, content)
+            self.assertIn("name=%s" % optional_repo_name, content)
+            self.assertIn("baseurl=http://example.com/repo/x86_64/optional", content)
+            self.assertIn("gpgcheck=0", content)
+
+        extra_repo_name = "extra-%s" % timestamp
+        extra_repo = os.path.join(configdir, "%s.repo" % extra_repo_name)
+        self.assertTrue(os.path.isfile(extra_repo))
+        with open(extra_repo, 'r') as f:
+            content = f.read()
+            self.assertIn("[%s]" % extra_repo_name, content)
+            self.assertIn("name=%s" % extra_repo_name, content)
+            self.assertIn("baseurl=http://example.com/repo/x86_64/extra", content)
+            self.assertIn("gpgcheck=0", content)
+
+        treeconf = json.load(open(treefile, 'r'))
+        repos = treeconf['repos']
+        self.assertEqual(len(repos), 3)
+        for name in [source_repo_from_name, optional_repo_name, extra_repo_name]:
+            self.assertIn(name, repos)
+
+    @mock.patch('pungi.ostree.utils.datetime')
+    @mock.patch('kobo.shortcuts.run')
+    def test_extra_config_with_keep_original_sources(self, run, time):
+        time.datetime.now.return_value = datetime.datetime(2016, 1, 1, 1, 1)
+        timestamp = time.datetime.now().strftime("%Y%m%d%H%M%S")
+
+        configdir = os.path.join(self.topdir, 'config')
+        self._make_dummy_config_dir(configdir)
+        treefile = os.path.join(configdir, 'fedora-atomic-docker-host.json')
+
+        repo = os.path.join(self.topdir, 'atomic')
+
+        extra_config_file = os.path.join(self.topdir, 'extra_config.json')
+        extra_config = {
+            "source_repo_from": "http://www.example.com/Server.repo",
+            "extra_source_repos": [
+                {
+                    "name": "optional",
+                    "baseurl": "http://example.com/repo/x86_64/optional",
+                    "exclude": "systemd-container",
+                    "gpgcheck": False
+                },
+                {
+                    "name": "extra",
+                    "baseurl": "http://example.com/repo/x86_64/extra",
+                }
+            ],
+            "keep_original_sources": True
+        }
+        helpers.touch(extra_config_file, json.dumps(extra_config))
+
+        ostree.main([
+            'tree',
+            '--repo=%s' % repo,
+            '--log-dir=%s' % os.path.join(self.topdir, 'logs', 'Atomic'),
+            '--treefile=%s' % treefile,
+            '--extra-config=%s' % extra_config_file,
+        ])
+
+        source_repo_from_name = "source_repo_from-%s" % timestamp
+        optional_repo_name = "optional-%s" % timestamp
+        extra_repo_name = "extra-%s" % timestamp
+
+        treeconf = json.load(open(treefile, 'r'))
+        repos = treeconf['repos']
+        self.assertEqual(len(repos), 6)
+        for name in ['fedora-rawhide', 'fedora-24', 'fedora-23',
+                     source_repo_from_name, optional_repo_name, extra_repo_name]:
+            self.assertIn(name, repos)
+
 if __name__ == '__main__':
     unittest.main()