lib/aggro/handler/query.rb in aggro-0.0.3 vs lib/aggro/handler/query.rb in aggro-0.0.4

@@ -35,26 +35,38 @@
       def handle_foreign
         comandee_locator.primary_node.client.post message
       end
 
       def handle_known
-        if channel.handles_query?(query)
-          result = channel.run_query(query)
-
-          Message::Result.new result.value(5)
+        if channel
+          if channel.handles_query?(query)
+            handle_supported
+          else
+            Message::UnhandledOperation.new
+          end
         else
-          Message::UnhandledOperation.new
+          Message::InvalidTarget.new
         end
-      rescue NoMethodError
-        Message::InvalidTarget.new
       end
 
       def handle_local
         query_known? ? handle_known : handle_unknown
       end
 
       def handle_unknown
         Message::UnknownOperation.new
+      end
+
+      def handle_supported
+        result = channel.run_query(query)
+
+        result.wait(5)
+
+        if result.fulfilled?
+          Message::Result.new result.value
+        else
+          Message::Result.new Aggro::QueryError.new('Query timed out')
+        end
       end
     end
   end
 end