From bc59fb2d765f697272bf030882da650cf72e3692 Mon Sep 17 00:00:00 2001
From: Ilya Shakhat <ishakhat@mirantis.com>
Date: Thu, 23 Apr 2015 13:10:27 +0300
Subject: [PATCH] Adopt shaker.lib to recent result model changes

Change-Id: I153e822d6a83a332a248df22a55aed97aa22edf1
---
 shaker/lib.py     | 20 +++++++++++++++-----
 tests/test_lib.py | 37 +++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+), 5 deletions(-)
 create mode 100644 tests/test_lib.py

diff --git a/shaker/lib.py b/shaker/lib.py
index eec7cf9..2c4c69b 100644
--- a/shaker/lib.py
+++ b/shaker/lib.py
@@ -25,8 +25,20 @@ LOG = logging.getLogger(__name__)
 class Shaker(object):
     """How to use Shaker as library
 
-    shaker = Shaker('127.0.0.1:5999', ['the-agent'])
-    res = shaker.run_program('the-agent', 'ls -al')
+>>> from shaker import lib
+>>> shaker = lib.Shaker('127.0.0.1:5999', ['the-agent'])
+>>> shaker.run_program('the-agent', 'date')
+{
+    'status': 'ok',
+    'stdout': 'Thu Apr 23 13:09:08 EAT 2015\n',
+    'agent': 'the-agent',
+    'command': {'data': 'date', 'type': 'program'},
+    'stderr': u'',
+    'executor': 'shell',
+    'test': 'shell',
+    'type': 'agent',
+    'id': '3a7c3d97-45f1-43ba-8460-2e37e679e3d5'
+}
     """
     def __init__(self, server_endpoint, agent_ids, polling_interval=1,
                  agent_loss_timeout=60, agent_join_timeout=600):
@@ -43,9 +55,7 @@ class Shaker(object):
         execution = {'tests': [test]}
         execution_result = server.execute(self.quorum, execution, agents)
 
-        results_per_iteration = execution_result[0]['results_per_iteration']
-        results_per_agent = results_per_iteration[0]['results_per_agent']
-        return dict((s['agent']['id'], s) for s in results_per_agent)
+        return list(execution_result.values())[0]
 
     def run_program(self, agent_id, program):
         return self._run(agent_id, {'program': program})
diff --git a/tests/test_lib.py b/tests/test_lib.py
new file mode 100644
index 0000000..16bd111
--- /dev/null
+++ b/tests/test_lib.py
@@ -0,0 +1,37 @@
+# Copyright (c) 2015 Mirantis Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+# implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import mock
+import testtools
+
+from shaker import lib
+
+
+class TestLib(testtools.TestCase):
+
+    @mock.patch('shaker.engine.quorum.make_quorum')
+    def test_run_program(self, make_quorum_patch):
+        quorum_mock = mock.MagicMock()
+        make_quorum_patch.return_value = quorum_mock
+
+        quorum_mock.execute = mock.Mock(
+            return_value={'AGENT': {'status': 'ok', 'stdout': 'STDOUT', }})
+
+        shaker = lib.Shaker('127.0.0.1:5999', ['AGENT'])
+        res = shaker.run_program('AGENT', 'ls -al')
+
+        self.assertDictContainsSubset(
+            {'status': 'ok', 'stdout': 'STDOUT', 'agent': 'AGENT',
+             'executor': 'shell', 'type': 'agent'}, res)