From c884d916e03a105f7c84e343fcd01be090dad1bb Mon Sep 17 00:00:00 2001
From: Masahiro Nakagawa <repeatedly@gmail.com>
Date: Mon, 21 May 2018 16:02:17 +0900
Subject: [PATCH] counter: Add wait API

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

Signed-off-by: Masahiro Nakagawa <repeatedly@gmail.com>
---
 lib/fluent/counter/client.rb | 11 ++++++++++-
 lib/fluent/counter/error.rb  | 21 +++++++++++++++++++++
 test/counter/test_client.rb  |  4 ++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/lib/fluent/counter/client.rb b/lib/fluent/counter/client.rb
index fc2ea9989a..1f426e5a23 100644
--- a/lib/fluent/counter/client.rb
+++ b/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
diff --git a/lib/fluent/counter/error.rb b/lib/fluent/counter/error.rb
index 1052d122c5..45b8877238 100644
--- a/lib/fluent/counter/error.rb
+++ b/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
diff --git a/test/counter/test_client.rb b/test/counter/test_client.rb
index bab55b5a1f..18fc0d65ca 100644
--- a/test/counter/test_client.rb
+++ b/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