Mercurial > hg > thermostat-ng > agent

changeset 1755:e5dc352ea8be

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Remove requirement for hostId option for dump-heap. Replace Service lookup in command with activator injecting services. Remove hostId argument for dump-heap command. Updating of relevant tests. PR2396 Reviewed-by: omajid Review-thread: http://icedtea.classpath.org/pipermail/thermostat/2015-June/014043.html
author James Aziz <jaziz@redhat.com>
date Fri, 12 Jun 2015 16:46:19 -0400
parents 821663c6872c
children 6e5f6b43160a
files killvm/command/src/main/java/com/redhat/thermostat/killvm/command/internal/Activator.java killvm/command/src/test/java/com/redhat/thermostat/killvm/command/KillVmCommandTest.java vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/Activator.java vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapCommand.java vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapHelper.java vm-heap-analysis/command/src/test/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapCommandTest.java vm-heap-analysis/command/src/test/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapHelperTest.java vm-heap-analysis/distribution/thermostat-plugin.xml
diffstat 8 files changed, 167 insertions(+), 121 deletions(-) [+]
line wrap: on
line diff
--- a/killvm/command/src/main/java/com/redhat/thermostat/killvm/command/internal/Activator.java	Fri Jun 12 10:58:21 2015 -0400
+++ b/killvm/command/src/main/java/com/redhat/thermostat/killvm/command/internal/Activator.java	Fri Jun 12 16:46:19 2015 -0400
@@ -71,7 +71,7 @@
             public void dependenciesAvailable(Map<String, Object> services) {
                 AgentInfoDAO agentDao = (AgentInfoDAO) services.get(AgentInfoDAO.class.getName());
                 VmInfoDAO vmDao = (VmInfoDAO) services.get(VmInfoDAO.class.getName());
-                KillVMRequest request= (KillVMRequest) services.get(KillVMRequest.class.getName());
+                KillVMRequest request = (KillVMRequest) services.get(KillVMRequest.class.getName());
 
                 command.setAgentInfoDAO(agentDao);
                 command.setVmInfoDAO(vmDao);
--- a/killvm/command/src/test/java/com/redhat/thermostat/killvm/command/KillVmCommandTest.java	Fri Jun 12 10:58:21 2015 -0400
+++ b/killvm/command/src/test/java/com/redhat/thermostat/killvm/command/KillVmCommandTest.java	Fri Jun 12 16:46:19 2015 -0400
@@ -93,9 +93,7 @@
         setServices();
         cmd.run(ctx);
 
-        verify(request).sendKillVMRequestToAgent(any(AgentId.class), any(int.class), any(AgentInfoDAO.class), any
-                (RequestResponseListener
-                .class));
+        verify(request).sendKillVMRequestToAgent(any(AgentId.class), any(int.class), any(AgentInfoDAO.class), any(RequestResponseListener.class));
     }
 
     @Test(expected = CommandException.class)
--- a/vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/Activator.java	Fri Jun 12 10:58:21 2015 -0400
+++ b/vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/Activator.java	Fri Jun 12 16:46:19 2015 -0400
@@ -36,28 +36,64 @@
 
 package com.redhat.thermostat.vm.heap.analysis.command.internal;
 
+import java.util.Map;
+
 import org.osgi.framework.BundleActivator;
 import org.osgi.framework.BundleContext;
 
+import com.redhat.thermostat.client.command.RequestQueue;
+import com.redhat.thermostat.common.MultipleServiceTracker;
 import com.redhat.thermostat.common.cli.Command;
 import com.redhat.thermostat.common.cli.CommandRegistry;
 import com.redhat.thermostat.common.cli.CommandRegistryImpl;
+import com.redhat.thermostat.storage.dao.AgentInfoDAO;
+import com.redhat.thermostat.storage.dao.VmInfoDAO;
 
 public class Activator implements BundleActivator {
 
     private CommandRegistry reg;
+    private MultipleServiceTracker serviceTracker;
+    private DumpHeapCommand dumpHeapCommand = new DumpHeapCommand();
 
     @Override
     public void start(final BundleContext context) throws Exception {
         reg = new CommandRegistryImpl(context);
 
-        registerCommand("dump-heap", new DumpHeapCommand());
+        registerCommand("dump-heap", dumpHeapCommand);
         registerCommand("list-heap-dumps", new ListHeapDumpsCommand());
         registerCommand("save-heap-dump-to-file", new SaveHeapDumpToFileCommand());
         registerCommand("show-heap-histogram", new ShowHeapHistogramCommand());
         registerCommand("find-objects", new FindObjectsCommand());
         registerCommand("object-info", new ObjectInfoCommand());
         registerCommand("find-root", new FindRootCommand());
+
+        Class<?>[] serviceDeps = new Class<?>[] {
+                AgentInfoDAO.class,
+                VmInfoDAO.class,
+                RequestQueue.class,
+        };
+
+        serviceTracker = new MultipleServiceTracker(context, serviceDeps, new MultipleServiceTracker.Action() {
+            @Override
+            public void dependenciesAvailable(Map<String, Object> services) {
+                VmInfoDAO vmDao = (VmInfoDAO) services.get(VmInfoDAO.class.getName());
+                AgentInfoDAO agentDao = (AgentInfoDAO) services.get(AgentInfoDAO.class.getName());
+                RequestQueue queue = (RequestQueue) services.get(RequestQueue.class.getName());
+
+                dumpHeapCommand.setAgentInfoDAO(agentDao);
+                dumpHeapCommand.setVmInfoDAO(vmDao);
+                dumpHeapCommand.setRequestQueue(queue);
+            }
+
+            @Override
+            public void dependenciesUnavailable() {
+                dumpHeapCommand.setAgentInfoDAO(null);
+                dumpHeapCommand.setVmInfoDAO(null);
+                dumpHeapCommand.setRequestQueue(null);
+            }
+        });
+
+        serviceTracker.open();
     }
 
     private void registerCommand(String name, Command command) {
@@ -66,6 +102,7 @@
 
     @Override
     public void stop(BundleContext context) throws Exception {
+        serviceTracker.close();
         reg.unregisterCommands();
     }
     
--- a/vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapCommand.java	Fri Jun 12 10:58:21 2015 -0400
+++ b/vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapCommand.java	Fri Jun 12 16:46:19 2015 -0400
@@ -38,18 +38,17 @@
 
 import java.util.concurrent.Semaphore;
 
-import org.osgi.framework.BundleContext;
-import org.osgi.framework.FrameworkUtil;
-import org.osgi.framework.ServiceReference;
-
 import com.redhat.thermostat.client.cli.HostVMArguments;
 import com.redhat.thermostat.client.command.RequestQueue;
 import com.redhat.thermostat.common.cli.AbstractCommand;
 import com.redhat.thermostat.common.cli.CommandContext;
 import com.redhat.thermostat.common.cli.CommandException;
 import com.redhat.thermostat.shared.locale.Translate;
+import com.redhat.thermostat.storage.core.AgentId;
+import com.redhat.thermostat.storage.core.VmId;
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.dao.VmInfoDAO;
+import com.redhat.thermostat.storage.model.VmInfo;
 import com.redhat.thermostat.vm.heap.analysis.command.locale.LocaleResources;
 
 
@@ -57,24 +56,35 @@
 
     private static final Translate<LocaleResources> translator = LocaleResources.createLocalizer();
 
-    private BundleContext context;
     private final DumpHeapHelper implementation;
 
+    private VmInfoDAO vmInfoDAO;
+    private AgentInfoDAO agentInfoDAO;
+    private RequestQueue queue;
+
     public DumpHeapCommand() {
-        this(FrameworkUtil.getBundle(DumpHeapCommand.class).getBundleContext(), new DumpHeapHelper());
+        this(new DumpHeapHelper());
     }
 
-    DumpHeapCommand(BundleContext context, DumpHeapHelper impl) {
-        this.context = context;
+    DumpHeapCommand(DumpHeapHelper impl) {
         this.implementation = impl;
     }
 
     @Override
     public void run(final CommandContext ctx) throws CommandException {
-        final HostVMArguments args = new HostVMArguments(ctx.getArguments());
+        requireNonNull(vmInfoDAO, translator.localize(LocaleResources.VM_SERVICE_UNAVAILABLE));
+        requireNonNull(agentInfoDAO, translator.localize(LocaleResources.AGENT_SERVICE_UNAVAILABLE));
+        requireNonNull(queue, translator.localize(LocaleResources.REQUEST_QUEUE_UNAVAILABLE));
+
+        final HostVMArguments args = new HostVMArguments(ctx.getArguments(), false, true);
+
+        VmId vmId = new VmId(args.getVM().getVmId());
+        final VmInfo vmInfo = vmInfoDAO.getVmInfo(vmId);
+        final AgentId agentId = new AgentId(vmInfo.getAgentId());
 
         final CommandException[] ex = new CommandException[1];
         final Semaphore s = new Semaphore(0);
+
         Runnable successHandler = new Runnable() {
             @Override
             public void run() {
@@ -82,32 +92,16 @@
                 s.release();
             }
         };
+
         Runnable errorHandler = new Runnable() {
             public void run() {
                 ex[0] = new CommandException(translator.localize(
-                        LocaleResources.HEAP_DUMP_ERROR, args.getHost()
-                                .getStringID(), args.getVM().getVmId()));
+                        LocaleResources.HEAP_DUMP_ERROR, vmInfo.getAgentId(), vmInfo.getVmId()));
                 s.release();
             }
         };
 
-        ServiceReference vmInfoRef = context.getServiceReference(VmInfoDAO.class.getName());
-        requireNonNull(vmInfoRef, translator.localize(LocaleResources.VM_SERVICE_UNAVAILABLE));
-        VmInfoDAO vmInfoDAO = (VmInfoDAO) context.getService(vmInfoRef);
-        
-        ServiceReference agentInfoRef = context.getServiceReference(AgentInfoDAO.class.getName());
-        requireNonNull(agentInfoRef, translator.localize(LocaleResources.AGENT_SERVICE_UNAVAILABLE));
-        AgentInfoDAO agentInfoDAO = (AgentInfoDAO) context.getService(agentInfoRef);
-        
-        ServiceReference requestQueueRef = context.getServiceReference(RequestQueue.class.getName());
-        requireNonNull(requestQueueRef, translator.localize(LocaleResources.REQUEST_QUEUE_UNAVAILABLE));
-        RequestQueue queue = (RequestQueue) context.getService(requestQueueRef);
-        
-        implementation.execute(vmInfoDAO, agentInfoDAO, args.getVM(), queue, successHandler, errorHandler);
-        
-        context.ungetService(vmInfoRef);
-        context.ungetService(agentInfoRef);
-        context.ungetService(requestQueueRef);
+        implementation.execute(vmInfoDAO, agentInfoDAO, agentId, vmId, queue, successHandler, errorHandler);
         
         try {
             s.acquire();
@@ -120,5 +114,17 @@
         }
     }
 
+    public void setVmInfoDAO(VmInfoDAO vmInfoDAO) {
+        this.vmInfoDAO = vmInfoDAO;
+    }
+
+    public void setAgentInfoDAO(AgentInfoDAO agentInfoDAO) {
+        this.agentInfoDAO = agentInfoDAO;
+    }
+
+    public void setRequestQueue(RequestQueue queue) {
+        this.queue = queue;
+    }
+
 }
 
--- a/vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapHelper.java	Fri Jun 12 10:58:21 2015 -0400
+++ b/vm-heap-analysis/command/src/main/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapHelper.java	Fri Jun 12 16:46:19 2015 -0400
@@ -44,8 +44,8 @@
 import com.redhat.thermostat.common.command.RequestResponseListener;
 import com.redhat.thermostat.common.command.Response;
 import com.redhat.thermostat.common.command.Response.ResponseType;
-import com.redhat.thermostat.storage.core.HostRef;
-import com.redhat.thermostat.storage.core.VmRef;
+import com.redhat.thermostat.storage.core.AgentId;
+import com.redhat.thermostat.storage.core.VmId;
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.dao.VmInfoDAO;
 import com.redhat.thermostat.storage.model.VmInfo;
@@ -82,24 +82,22 @@
 
     }
 
-    public void execute(VmInfoDAO vmInfoDAO, AgentInfoDAO agentInfoDAO, VmRef reference,
-            RequestQueue queue, Runnable heapDumpSuccessAction,
-            Runnable heapDumpFailureAction) {
+    public void execute(VmInfoDAO vmInfoDAO, AgentInfoDAO agentInfoDAO, AgentId agentId, VmId vmId,
+                        RequestQueue queue, Runnable heapDumpSuccessAction,
+                        Runnable heapDumpFailureAction) {
         // Get PID
-        VmInfo info = vmInfoDAO.getVmInfo(reference);
+        VmInfo info = vmInfoDAO.getVmInfo(vmId);
         int pid = info.getVmPid();
         
-        HostRef targetHostRef = reference.getHostRef();
-        InetSocketAddress target = agentInfoDAO.getAgentInformation(targetHostRef).getRequestQueueAddress();
+        InetSocketAddress target = agentInfoDAO.getAgentInformation(agentId).getRequestQueueAddress();
         Request req = new Request(RequestType.RESPONSE_EXPECTED, target);
         req.setReceiver(RECEIVER_CLASS_NAME);
         req.setParameter(Request.ACTION, CMD_CHANNEL_ACTION_NAME);
-        req.setParameter(VM_ID_PARAM, reference.getVmId());
+        req.setParameter(VM_ID_PARAM, vmId.get());
         req.setParameter(VM_PID_PARAM, String.valueOf(pid));
         req.addListener(new HeapDumpListener(heapDumpSuccessAction, heapDumpFailureAction));
 
         queue.putRequest(req);
-
     }
 
 }
--- a/vm-heap-analysis/command/src/test/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapCommandTest.java	Fri Jun 12 10:58:21 2015 -0400
+++ b/vm-heap-analysis/command/src/test/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapCommandTest.java	Fri Jun 12 16:46:19 2015 -0400
@@ -45,6 +45,7 @@
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 import org.junit.Test;
 import org.mockito.ArgumentCaptor;
@@ -55,11 +56,12 @@
 import com.redhat.thermostat.common.cli.CommandException;
 import com.redhat.thermostat.common.cli.SimpleArguments;
 import com.redhat.thermostat.shared.locale.Translate;
-import com.redhat.thermostat.storage.core.VmRef;
+import com.redhat.thermostat.storage.core.AgentId;
+import com.redhat.thermostat.storage.core.VmId;
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.dao.VmInfoDAO;
+import com.redhat.thermostat.storage.model.VmInfo;
 import com.redhat.thermostat.test.TestCommandContextFactory;
-import com.redhat.thermostat.testutils.StubBundleContext;
 import com.redhat.thermostat.vm.heap.analysis.command.locale.LocaleResources;
 
 public class DumpHeapCommandTest {
@@ -68,14 +70,13 @@
             .createLocalizer();
 
     @Test
-    public void verifyAcuallyCallsWorker() throws CommandException {
+    public void verifyActuallyCallsWorker() throws CommandException {
         VmInfoDAO vmInfoDao = mock(VmInfoDAO.class);
         AgentInfoDAO agentInfoDao = mock(AgentInfoDAO.class);
         RequestQueue queue = mock(RequestQueue.class);
-        StubBundleContext context = new StubBundleContext();
-        context.registerService(VmInfoDAO.class, vmInfoDao, null);
-        context.registerService(AgentInfoDAO.class, agentInfoDao, null);
-        context.registerService(RequestQueue.class, queue, null);
+        VmInfo vmInfo = new VmInfo("myAgent", "foo", 123, 0, 0, null, null, null, null, null, null, null, null, null,
+                null, null,0, "myUsername");
+        when(vmInfoDao.getVmInfo(new VmId("foo"))).thenReturn(vmInfo);
 
         DumpHeapHelper impl = mock(DumpHeapHelper.class);
         final ArgumentCaptor<Runnable> successHandler = ArgumentCaptor
@@ -87,36 +88,39 @@
                 successHandler.getValue().run();
                 return null;
             }
-        }).when(impl).execute(eq(vmInfoDao), eq(agentInfoDao), any(VmRef.class), eq(queue),
+        }).when(impl).execute(eq(vmInfoDao), eq(agentInfoDao), any(AgentId.class), any(VmId.class), eq(queue),
                 successHandler.capture(), any(Runnable.class));
 
-        DumpHeapCommand command = new DumpHeapCommand(context, impl);
+        DumpHeapCommand command = new DumpHeapCommand(impl);
+
+        command.setVmInfoDAO(vmInfoDao);
+        command.setAgentInfoDAO(agentInfoDao);
+        command.setRequestQueue(queue);
 
         TestCommandContextFactory factory = new TestCommandContextFactory();
 
         SimpleArguments args = new SimpleArguments();
-        args.addArgument("hostId", "foo");
-        args.addArgument("vmId", "bar");
+        args.addArgument("vmId", "foo");
 
         command.run(factory.createContext(args));
 
-        verify(impl).execute(eq(vmInfoDao), eq(agentInfoDao), isA(VmRef.class), eq(queue),
+        verify(impl).execute(eq(vmInfoDao), eq(agentInfoDao), isA(AgentId.class), isA(VmId.class), eq(queue),
                 any(Runnable.class), any(Runnable.class));
         assertEquals("Done\n", factory.getOutput());
     }
 
     @Test
-    public void verifyNeedsHostAndVmId() throws CommandException {
+    public void verifyNeedsVmId() throws CommandException {
         VmInfoDAO vmInfoDao = mock(VmInfoDAO.class);
         AgentInfoDAO agentInfoDao = mock(AgentInfoDAO.class);
         RequestQueue queue = mock(RequestQueue.class);
-        StubBundleContext context = new StubBundleContext();
-        context.registerService(VmInfoDAO.class, vmInfoDao, null);
-        context.registerService(AgentInfoDAO.class, agentInfoDao, null);
-        context.registerService(RequestQueue.class, queue, null);
 
         DumpHeapHelper impl = mock(DumpHeapHelper.class);
-        DumpHeapCommand command = new DumpHeapCommand(context, impl);
+        DumpHeapCommand command = new DumpHeapCommand(impl);
+
+        command.setVmInfoDAO(vmInfoDao);
+        command.setAgentInfoDAO(agentInfoDao);
+        command.setRequestQueue(queue);
 
         TestCommandContextFactory factory = new TestCommandContextFactory();
 
@@ -126,7 +130,7 @@
             command.run(factory.createContext(args));
             assertTrue("should not reach here", false);
         } catch (CommandException ce) {
-            assertEquals("a hostId is required", ce.getMessage());
+            assertEquals("a vmId is required", ce.getMessage());
         }
     }
 
@@ -134,17 +138,16 @@
     public void verifyFailsIfAgentDaoIsNotAvailable() {
         VmInfoDAO vmInfoDao = mock(VmInfoDAO.class);
         RequestQueue queue = mock(RequestQueue.class);
-        StubBundleContext context = new StubBundleContext();
-        context.registerService(VmInfoDAO.class, vmInfoDao, null);
-        context.registerService(RequestQueue.class, queue, null);
 
         DumpHeapHelper impl = mock(DumpHeapHelper.class);
-        DumpHeapCommand command = new DumpHeapCommand(context, impl);
+        DumpHeapCommand command = new DumpHeapCommand(impl);
+
+        command.setVmInfoDAO(vmInfoDao);
+        command.setRequestQueue(queue);
 
         TestCommandContextFactory factory = new TestCommandContextFactory();
 
         SimpleArguments args = new SimpleArguments();
-        args.addArgument("hostId", "foo");
         args.addArgument("vmId", "bar");
 
         try {
@@ -159,17 +162,16 @@
     public void verifyFailsIfRequestQueueIsNotAvailable() {
         VmInfoDAO vmInfoDao = mock(VmInfoDAO.class);
         AgentInfoDAO agentInfoDao = mock(AgentInfoDAO.class);
-        StubBundleContext context = new StubBundleContext();
-        context.registerService(VmInfoDAO.class, vmInfoDao, null);
-        context.registerService(AgentInfoDAO.class, agentInfoDao, null);
 
         DumpHeapHelper impl = mock(DumpHeapHelper.class);
-        DumpHeapCommand command = new DumpHeapCommand(context, impl);
+        DumpHeapCommand command = new DumpHeapCommand(impl);
+
+        command.setVmInfoDAO(vmInfoDao);
+        command.setAgentInfoDAO(agentInfoDao);
 
         TestCommandContextFactory factory = new TestCommandContextFactory();
 
         SimpleArguments args = new SimpleArguments();
-        args.addArgument("hostId", "foo");
         args.addArgument("vmId", "bar");
         args.addArgument("vmPid", "123");
 
@@ -185,17 +187,16 @@
     public void verifyFailsIfVmDaoIsNotAvailable() {
         AgentInfoDAO agentInfoDao = mock(AgentInfoDAO.class);
         RequestQueue queue = mock(RequestQueue.class);
-        StubBundleContext context = new StubBundleContext();
-        context.registerService(AgentInfoDAO.class, agentInfoDao, null);
-        context.registerService(RequestQueue.class, queue, null);
 
         DumpHeapHelper impl = mock(DumpHeapHelper.class);
-        DumpHeapCommand command = new DumpHeapCommand(context, impl);
+        DumpHeapCommand command = new DumpHeapCommand(impl);
+
+        command.setAgentInfoDAO(agentInfoDao);
+        command.setRequestQueue(queue);
 
         TestCommandContextFactory factory = new TestCommandContextFactory();
 
         SimpleArguments args = new SimpleArguments();
-        args.addArgument("hostId", "foo");
         args.addArgument("vmId", "bar");
 
         try {
@@ -208,41 +209,52 @@
 
     @Test
     public void verifyErrorMessage() {
-        final String HOST_ID = "myHost";
+        final String AGENT_ID = "myAgent";
         final String VM_ID = "myVm";
         VmInfoDAO vmInfoDao = mock(VmInfoDAO.class);
         AgentInfoDAO agentInfoDao = mock(AgentInfoDAO.class);
         RequestQueue queue = mock(RequestQueue.class);
-        StubBundleContext context = new StubBundleContext();
-        context.registerService(VmInfoDAO.class, vmInfoDao, null);
-        context.registerService(AgentInfoDAO.class, agentInfoDao, null);
-        context.registerService(RequestQueue.class, queue, null);
+        VmInfo vmInfo = new VmInfo(AGENT_ID, VM_ID, 123, 0, 0, null, null, null, null, null, null, null, null,
+                null, null, null, 0, null);
+        VmId vmId = new VmId(VM_ID);
+        AgentId agentId = new AgentId(AGENT_ID);
+
+        when(vmInfoDao.getVmInfo(vmId)).thenReturn(vmInfo);
 
         DumpHeapHelper impl = mock(DumpHeapHelper.class);
-        DumpHeapCommand command = new DumpHeapCommand(context, impl);
+        DumpHeapCommand command = new DumpHeapCommand(impl);
+
+        final ArgumentCaptor<Runnable> errorHandler = ArgumentCaptor.forClass(Runnable.class);
+
+        doAnswer(new Answer<Void>() {
+            @Override
+            public Void answer(InvocationOnMock invocation) throws Throwable {
+                errorHandler.getValue().run();
+                return null;
+            }
+        }).when(impl).execute(eq(vmInfoDao),
+                              eq(agentInfoDao),
+                              eq(agentId),
+                              eq(vmId),
+                              eq(queue),
+                              any(Runnable.class),
+                              errorHandler.capture());
+
+        command.setVmInfoDAO(vmInfoDao);
+        command.setAgentInfoDAO(agentInfoDao);
+        command.setRequestQueue(queue);
+
         TestCommandContextFactory factory = new TestCommandContextFactory();
 
         SimpleArguments args = new SimpleArguments();
-        args.addArgument("hostId", HOST_ID);
         args.addArgument("vmId", VM_ID);
 
-        doAnswer(new Answer<Object>() {
-
-            @Override
-            public Object answer(InvocationOnMock invocation) throws Throwable {
-                Runnable failRunnable = (Runnable) invocation.getArguments()[5];
-                failRunnable.run();
-                return null;
-            }
-        }).when(impl).execute(any(VmInfoDAO.class), any(AgentInfoDAO.class), any(VmRef.class), 
-                any(RequestQueue.class), any(Runnable.class), any(Runnable.class));
-
         try {
             command.run(factory.createContext(args));
             fail("CommandException expected");
         } catch (CommandException e) {
             assertEquals(TRANSLATOR.localize(LocaleResources.HEAP_DUMP_ERROR,
-                    HOST_ID, VM_ID).getContents(), e.getMessage());
+                    AGENT_ID, VM_ID).getContents(), e.getMessage());
         }
     }
 
--- a/vm-heap-analysis/command/src/test/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapHelperTest.java	Fri Jun 12 10:58:21 2015 -0400
+++ b/vm-heap-analysis/command/src/test/java/com/redhat/thermostat/vm/heap/analysis/command/internal/DumpHeapHelperTest.java	Fri Jun 12 16:46:19 2015 -0400
@@ -56,8 +56,8 @@
 import com.redhat.thermostat.common.command.RequestResponseListener;
 import com.redhat.thermostat.common.command.Response;
 import com.redhat.thermostat.common.command.Response.ResponseType;
-import com.redhat.thermostat.storage.core.HostRef;
-import com.redhat.thermostat.storage.core.VmRef;
+import com.redhat.thermostat.storage.core.AgentId;
+import com.redhat.thermostat.storage.core.VmId;
 import com.redhat.thermostat.storage.dao.AgentInfoDAO;
 import com.redhat.thermostat.storage.dao.VmInfoDAO;
 import com.redhat.thermostat.storage.model.AgentInformation;
@@ -68,7 +68,9 @@
     private VmInfoDAO vmInfoDAO;
     private AgentInfoDAO agentInfoDao;
     private DumpHeapHelper cmd;
-    private VmRef vmRef;
+    private VmId vmId;
+    private AgentId agentId;
+    private AgentInformation agentInfo;
     private RequestQueue reqQueue;
     private Runnable heapDumpCompleteAction;
     private Runnable heapDumpFailedAction;
@@ -77,23 +79,24 @@
     public void setUp() {
         reqQueue = mock(RequestQueue.class);
 
-        HostRef host = mock(HostRef.class);
-
-        AgentInformation agentInfo = mock(AgentInformation.class);
-        when(agentInfo.getRequestQueueAddress()).thenReturn(new InetSocketAddress("test", 123));
-
         agentInfoDao = mock(AgentInfoDAO.class);
-        when(agentInfoDao.getAgentInformation(host)).thenReturn(agentInfo);
 
         cmd = new DumpHeapHelper();
-        vmRef = mock(VmRef.class);
-        when(vmRef.getVmId()).thenReturn("vmId");
-        when(vmRef.getHostRef()).thenReturn(host);
-        
+
+        agentInfo = mock(AgentInformation.class);
+        when(agentInfo.getRequestQueueAddress()).thenReturn(new InetSocketAddress("test", 123));
+
+        agentId = mock(AgentId.class);
+        when(agentInfoDao.getAgentInformation(agentId)).thenReturn(agentInfo);
+
+        vmId = mock(VmId.class);
+        when(vmId.get()).thenReturn("vmId");
+
         VmInfo vmInfo = mock(VmInfo.class);
         when(vmInfo.getVmPid()).thenReturn(123);
+
         vmInfoDAO = mock(VmInfoDAO.class);
-        when(vmInfoDAO.getVmInfo(vmRef)).thenReturn(vmInfo);
+        when(vmInfoDAO.getVmInfo(vmId)).thenReturn(vmInfo);
         
         heapDumpCompleteAction = mock(Runnable.class);
         heapDumpFailedAction = mock(Runnable.class);
@@ -102,15 +105,14 @@
     @After
     public void tearDown() {
         heapDumpCompleteAction = null;
-        vmRef = null;
+        vmId = null;
         cmd = null;
         reqQueue = null;
     }
 
     @Test
     public void testExecute() {
-
-        cmd.execute(vmInfoDAO, agentInfoDao, vmRef, reqQueue, heapDumpCompleteAction, heapDumpFailedAction);
+        cmd.execute(vmInfoDAO, agentInfoDao, agentId, vmId, reqQueue, heapDumpCompleteAction, heapDumpFailedAction);
 
         ArgumentCaptor<Request> reqArg = ArgumentCaptor.forClass(Request.class);
         verify(reqQueue).putRequest(reqArg.capture());
@@ -133,7 +135,7 @@
     @Test
     public void testExecuteFailure() {
 
-        cmd.execute(vmInfoDAO, agentInfoDao, vmRef, reqQueue, heapDumpCompleteAction, heapDumpFailedAction);
+        cmd.execute(vmInfoDAO, agentInfoDao, agentId, vmId, reqQueue, heapDumpCompleteAction, heapDumpFailedAction);
 
         ArgumentCaptor<Request> reqArg = ArgumentCaptor.forClass(Request.class);
         verify(reqQueue).putRequest(reqArg.capture());
--- a/vm-heap-analysis/distribution/thermostat-plugin.xml	Fri Jun 12 10:58:21 2015 -0400
+++ b/vm-heap-analysis/distribution/thermostat-plugin.xml	Fri Jun 12 16:46:19 2015 -0400
@@ -46,13 +46,6 @@
       <description>Trigger a heap dump on the VM.</description>
       <options>
         <option>
-          <long>hostId</long>
-          <short>a</short>
-          <argument>host</argument>
-          <required>true</required>
-          <description>the ID of the host to monitor</description>
-        </option>
-        <option>
           <long>vmId</long>
           <short>v</short>
           <argument>vm</argument>