counter: Add wait API

Similar to "get" but raises an error when API call has a problem.

Signed-off-by: Masahiro Nakagawa <[email protected]>

lib/fluent/counter/client.rb

@@ -16,6 +16,7 @@
 
 require 'cool.io'
 require 'fluent/counter/base_socket'
+require 'fluent/counter/error'
 require 'timeout'
 
 module Fluent
@@ -59,7 +60,7 @@ def stop
       def establish(scope)
         scope = Timeout.timeout(@timeout) {
           response = send_request('establish', nil, [scope])
-          raise response.errors.first if response.errors?
+          Fluent::Counter.raise_error(response.errors.first) if response.errors?
           data = response.data
           data.first
         }
@@ -274,6 +275,14 @@ def get
         @result
       end
 
+      def wait
+        res = get
+        if res.error?
+          Fluent::Counter.raise_error(res.errors.first)
+        end
+        res
+      end
+
       private
 
       def join

lib/fluent/counter/error.rb

@@ -61,5 +61,26 @@ def code
         'internal_server_error'
       end
     end
+
+    def raise_error(response)
+      msg = response['message']
+      case response['code']
+      when 'invalid_params'
+        raise InvalidParams.new(msg)
+      when 'unknown_key'
+        raise UnknownKey.new(msg)
+      when 'parse_error'
+        raise ParseError.new(msg)
+      when 'invalid_request'
+        raise InvalidRequest.new(msg)
+      when 'method_not_found'
+        raise MethodNotFound.new(msg)
+      when 'internal_server_error'
+        raise InternalServerError.new(msg)
+      else
+        raise "Unknown code: #{response['code']}"
+      end
+    end
+    module_function :raise_error
   end
 end

test/counter/test_client.rb

@@ -151,6 +151,10 @@ def extract_value_from_server(server, scope, name)
 
       assert_empty response.data
       assert_equal expected_error, errors
+
+      assert_raise {
+        @client.init(param).wait
+      }
     end
 
     test 'return an existing value when passed key already exists and ignore option is true' do