Regression: single custom exception not handled

9bebe7f missed updating the logic for handling a single custom exception (where the thrift id can be inferred).

swift-service/src/main/java/com/facebook/swift/service/metadata/ThriftMethodMetadata.java

@@ -34,10 +34,13 @@
 import java.lang.reflect.Method;
 import java.lang.reflect.Modifier;
 import java.lang.reflect.Type;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 
 import static com.facebook.swift.codec.metadata.ReflectionHelper.extractParameterNames;
+import static com.google.common.base.Preconditions.checkState;
 
 @Immutable
 public class ThriftMethodMetadata
@@ -156,18 +159,29 @@ public boolean getOneway() {
     private ImmutableMap<Short, ThriftType> buildExceptionMap(ThriftCatalog catalog,
                                                               ThriftMethod thriftMethod) {
         ImmutableMap.Builder<Short, ThriftType> exceptions = ImmutableMap.builder();
+        Set<Type> exceptionTypes = new HashSet<>();
+        int customExceptionCount = 0;
 
         if (thriftMethod.exception().length > 0) {
             for (ThriftException thriftException : thriftMethod.exception()) {
                 exceptions.put(thriftException.id(), catalog.getThriftType(thriftException.type()));
+                checkState(exceptionTypes.add(thriftException.type()), "ThriftMethod.exception contains more than one value for %s", thriftException.type());
             }
-        } else if (method.getExceptionTypes().length == 2) {
-            // Catch the case where the method declares exactly TWO thrown types: one must
-            // be the generic TException or an ancestor, so if the other is annotated as a
-            // @ThriftStruct then we infer a thrift declaration for that remaining exception
-            // type
-            for (Class<?> exceptionClass : method.getExceptionTypes()) {
-                if (exceptionClass.isAnnotationPresent(ThriftStruct.class)) {
+        }
+
+        for (Class<?> exceptionClass : method.getExceptionTypes()) {
+            if (exceptionClass.isAssignableFrom(TException.class)) {
+                // the built-in exception types don't need special treatment
+                continue;
+            }
+
+            if (exceptionClass.isAnnotationPresent(ThriftStruct.class)) {
+                ++customExceptionCount;
+
+                if (!exceptionTypes.contains(exceptionClass)) {
+                    // there is no rhyme or reason to the order exception types are given to us,
+                    // so we can only infer the id once
+                    checkState(customExceptionCount <= 1, "ThriftMethod.exception annotation value must be specified when more than one custom exception is thrown.");
                     exceptions.put((short) 1, catalog.getThriftType(exceptionClass));
                 }
             }

swift-service/src/test/java/com/facebook/swift/service/metadata/TestThriftMethodMetadata.java

@@ -0,0 +1,137 @@
+/**
+ * Copyright 2012 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may
+ * not use this file except in compliance with the License. You may obtain
+ * a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+package com.facebook.swift.service.metadata;
+
+import com.facebook.swift.codec.ThriftStruct;
+import com.facebook.swift.codec.metadata.ThriftCatalog;
+import com.facebook.swift.codec.metadata.ThriftType;
+import com.facebook.swift.service.ThriftException;
+import com.facebook.swift.service.ThriftMethod;
+import org.apache.thrift.TException;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.lang.reflect.Type;
+import java.util.Map;
+import java.util.TreeMap;
+
+// TODO add more tests... (currently only tests exception metadata)
+public class TestThriftMethodMetadata
+{
+    @Test
+    public void testNoExceptions() throws Exception
+    {
+        assertExceptions("noExceptions");
+    }
+
+    @Test
+    public void testAnnotatedExceptions() throws Exception
+    {
+        assertExceptions("annotatedExceptions", ExceptionA.class, ExceptionB.class);
+    }
+
+    @Test
+    public void testInferredException() throws Exception
+    {
+        assertExceptions("inferredException", ExceptionA.class);
+    }
+
+    @Test
+    public void testInferredExceptionWithTException() throws Exception
+    {
+        assertExceptions("inferredExceptionWithTException", ExceptionA.class);
+    }
+
+    @Test
+    public void testInferredExceptionWithRuntimeException() throws Exception
+    {
+        assertExceptions("inferredExceptionWithRuntimeException", ExceptionA.class);
+    }
+
+    @Test
+    public void testInferredExceptionWithRuntimeAndTException() throws Exception
+    {
+        assertExceptions("inferredExceptionWithRuntimeAndTException", ExceptionA.class);
+    }
+
+    @Test(expectedExceptions = IllegalStateException.class,
+          expectedExceptionsMessageRegExp = "ThriftMethod.exception annotation value must be specified when more than one custom exception is thrown.")
+    public void testUninferrableException() throws Exception
+    {
+        assertExceptions("uninferrableException");
+    }
+
+    private void assertExceptions(String methodName, Class<? extends Exception>... expectedExceptions) throws Exception
+    {
+        ThriftMethodMetadata metadata = new ThriftMethodMetadata(DummyService.class.getMethod(methodName), new ThriftCatalog());
+        Map<Short, Type> actualIdMap = new TreeMap<>();
+        Map<Short, Type> expectedIdMap = new TreeMap<>();
+
+        for (Map.Entry<Short, ThriftType> entry : metadata.getExceptions().entrySet()) {
+            actualIdMap.put(entry.getKey(), entry.getValue().getJavaType());
+        }
+
+        short expectedId = 1;
+
+        for (Class<? extends Exception> expectedException : expectedExceptions) {
+            expectedIdMap.put(expectedId, expectedException);
+            ++expectedId;
+        }
+
+        // string comparison produces more useful failure message (and is safe, given the types)
+        Assert.assertEquals(actualIdMap.toString(), expectedIdMap.toString());
+    }
+
+    @SuppressWarnings("unused")
+    public static interface DummyService
+    {
+        @ThriftMethod
+        public void noExceptions();
+
+        @ThriftMethod(
+                exception = {
+                        @ThriftException(id = 1, type = ExceptionA.class),
+                        @ThriftException(id = 2, type = ExceptionB.class)
+                }
+        )
+        public void annotatedExceptions() throws ExceptionA, ExceptionB;
+
+        @ThriftMethod
+        public void inferredException() throws ExceptionA;
+
+        @ThriftMethod
+        public void inferredExceptionWithTException() throws ExceptionA, TException;
+
+        @ThriftMethod
+        public void inferredExceptionWithRuntimeException() throws IllegalArgumentException, ExceptionA;
+
+        @ThriftMethod
+        public void inferredExceptionWithRuntimeAndTException() throws IllegalArgumentException, ExceptionA, TException;
+
+        @ThriftMethod(exception = {@ThriftException(id = 1, type = ExceptionA.class)})
+        public void uninferrableException() throws ExceptionA, ExceptionB;
+    }
+
+    @ThriftStruct
+    public static class ExceptionA extends Exception
+    {
+    }
+
+    @ThriftStruct
+    public static class ExceptionB extends Exception
+    {
+    }
+}

swift-service/src/test/java/com/facebook/swift/service/puma/swift/PumaReadService.java

@@ -28,25 +28,25 @@ public interface PumaReadService extends Closeable
 {
     @ThriftMethod
     List<ReadResultQueryInfo> getResult(List<ReadQueryInfoTimeString> reader)
-            throws TException, ReadSemanticException;
+            throws ReadSemanticException;
 
     @ThriftMethod
     List<ReadResultQueryInfoTimeString> getResultTimeString(List<ReadQueryInfoTimeString> reader)
-            throws TException, ReadSemanticException;
+            throws ReadSemanticException;
 
     @ThriftMethod
     List<ReadResultQueryInfo> mergeQueryAggregation(
             MergeAggregationQueryInfo mergeAggregationQueryInfo
     )
-            throws TException, ReadSemanticException;
+            throws ReadSemanticException;
 
     @ThriftMethod
     long latestQueryableTime(String category, String appName, List<Integer> bucketNumbers)
-            throws TException, ReadSemanticException;
+            throws ReadSemanticException;
 
     @ThriftMethod
     List<Long> latestQueryableTimes(String category, String appName, List<Integer> bucketNumbers)
-            throws TException, ReadSemanticException;
+            throws ReadSemanticException;
 
     @Override
     void close();