1 files changed, 51 insertions, 0 deletions
diff --git a/python/mozperftest/mozperftest/tests/test_profile.py b/python/mozperftest/mozperftest/tests/test_profile.py
new file mode 100644
index 0000000000..b110773566
--- /dev/null
+++ b/python/mozperftest/mozperftest/tests/test_profile.py
@@ -0,0 +1,51 @@
+#!/usr/bin/env python
+# This Source Code Form is subject to the terms of the Mozilla Public
+# License, v. 2.0. If a copy of the MPL was not distributed with this
+# file, You can obtain one at http://mozilla.org/MPL/2.0/.
+import os
+import mozunit
+from unittest import mock
+import tempfile
+
+from mozperftest.tests.support import get_running_env
+from mozperftest.system.profile import Profile, ProfileNotFoundError
+
+
+def test_profile():
+    mach_cmd, metadata, env = get_running_env()
+
+    with Profile(env, mach_cmd) as profile:
+        profile(metadata)
+        profile_dir = env.get_arg("profile-directory")
+        assert os.path.exists(profile_dir)
+
+    assert not os.path.exists(profile_dir)
+
+
+CALLS = [0]
+
+
+def _return_profile(*args, **kw):
+    if CALLS[0] == 0:
+        CALLS[0] = 1
+        raise ProfileNotFoundError()
+
+    tempdir = tempfile.mkdtemp()
+
+    return tempdir
+
+
+@mock.patch("mozperftest.system.profile.get_profile", new=_return_profile)
+def test_conditionedprofile():
+    mach_cmd, metadata, env = get_running_env(profile_conditioned=True)
+
+    with Profile(env, mach_cmd) as profile:
+        profile(metadata)
+        profile_dir = env.get_arg("profile-directory")
+        assert os.path.exists(profile_dir)
+
+    assert not os.path.exists(profile_dir)
+
+
+if __name__ == "__main__":
+    mozunit.main()