diff --git a/src/main/java/net/onrc/onos/util/GraphDBOperation.java b/src/main/java/net/onrc/onos/util/GraphDBOperation.java
index bcaa3af..8e6a74d 100644
--- a/src/main/java/net/onrc/onos/util/GraphDBOperation.java
+++ b/src/main/java/net/onrc/onos/util/GraphDBOperation.java
@@ -29,6 +29,7 @@
 	public ISwitchObject newSwitch() {
 		FramedGraph<TitanGraph> fg = conn.getFramedGraph();	
 		ISwitchObject obj = fg.addVertex(null,ISwitchObject.class);
+		if (obj != null) obj.setType("switch");
 		return obj;
 	}
 
@@ -90,6 +91,7 @@
 	public IPortObject newPort() {
 		FramedGraph<TitanGraph> fg = conn.getFramedGraph();	
 		IPortObject obj = fg.addVertex(null,IPortObject.class);
+		if (obj != null) obj.setType("port");
 		return obj;
 	}
 	
@@ -97,6 +99,7 @@
 	public IDeviceObject newDevice() {
 		FramedGraph<TitanGraph> fg = conn.getFramedGraph();	
 		IDeviceObject obj = fg.addVertex(null,IDeviceObject.class);
+		if (obj != null) obj.setType("device");
 		return obj;
 	}
 	
@@ -132,6 +135,7 @@
 	public IFlowPath newFlowPath() {
 		FramedGraph<TitanGraph> fg = conn.getFramedGraph();	
 		IFlowPath flowPath = fg.addVertex(null, IFlowPath.class);
+		if (flowPath != null) flowPath.setType("flow");
 		return flowPath;
 	}
 
@@ -179,6 +183,7 @@
 	public IFlowEntry newFlowEntry() {
 		FramedGraph<TitanGraph> fg = conn.getFramedGraph();	
 		IFlowEntry flowEntry = fg.addVertex(null, IFlowEntry.class);
+		if (flowEntry != null) flowEntry.setType("flow_entry");
 		return flowEntry;
 	}
 
diff --git a/src/test/java/net/onrc/onos/util/GraphDBOperationTest.java b/src/test/java/net/onrc/onos/util/GraphDBOperationTest.java
index f6e2a81..04ef815 100644
--- a/src/test/java/net/onrc/onos/util/GraphDBOperationTest.java
+++ b/src/test/java/net/onrc/onos/util/GraphDBOperationTest.java
@@ -73,6 +73,7 @@
 	@After
 	public void tearDown() throws Exception {
 		conn.close();
+		titanGraph.shutdown();
 	}
 
 	/**
@@ -86,7 +87,6 @@
 		assertFalse(switches.iterator().hasNext());
 
 		ISwitchObject sw = op.newSwitch();
-		sw.setType("switch");
 		sw.setDPID("123");
 		sw.setState(SwitchState.ACTIVE.toString());
 		conn.endTx(Transaction.COMMIT);
@@ -109,20 +109,16 @@
 		// make sure there is no switch
 		switches = op.getAllSwitches();
 		assertFalse(switches.iterator().hasNext());
-
 		
 		ISwitchObject sw = op.newSwitch();
-		sw.setType("switch");
 		sw.setDPID("123");
 		sw.setState(SwitchState.ACTIVE.toString());
 		conn.endTx(Transaction.COMMIT);
-
-		switches = op.getAllSwitches();
-		assertTrue(switches.iterator().hasNext());
 		
-		ISwitchObject obtained_sw = switches.iterator().next();
-		String obtained_dpid = obtained_sw.getDPID(); 
-		assertEquals("123", obtained_dpid);
+		sw = op.searchSwitch("123");
+		op.removeSwitch(sw);
+
+		assertNull(op.searchSwitch("123"));
 	}
 
 	/**