clean up

diff --git a/pkgs/swift2objc/test/integration/subscripts_input.swift b/pkgs/swift2objc/test/integration/subscripts_input.swift
index 7cfb34b..fa94356 100644
--- a/pkgs/swift2objc/test/integration/subscripts_input.swift
+++ b/pkgs/swift2objc/test/integration/subscripts_input.swift

@@ -1,4 +1,4 @@
-
+import Foundation
 public class SubscriptClass {
     public var stored: String = ""
 
@@ -36,16 +36,16 @@
     }
 }
 
-public class NonTrivialArgsSubscript {
-    public subscript(other: SubscriptClass) -> Int {
+public class NonTrivialTypesSubscript {
+    public subscript(other: SubscriptClass) -> SubscriptClass {
         get {
-            return other.stored.count
+            return other
         }
     }
 }
 
-public class OptionalSubscript {
-    public subscript(opt: Int?) -> String? {
+@objc public class OptionalSubscript: NSObject {
+    @objc public subscript(opt: String?) -> String? {
         get {
             return opt == nil ? nil : "Value: \(opt!)"
         }

diff --git a/pkgs/swift2objc/test/integration/subscripts_output.swift b/pkgs/swift2objc/test/integration/subscripts_output.swift
index 3a3b765..fb44dd8 100644
--- a/pkgs/swift2objc/test/integration/subscripts_output.swift
+++ b/pkgs/swift2objc/test/integration/subscripts_output.swift

@@ -2,6 +2,20 @@
 
 import Foundation
 
+@objc public class OptionalSubscriptWrapper: NSObject {
+  var wrappedInstance: OptionalSubscript
+
+  init(_ wrappedInstance: OptionalSubscript) {
+    self.wrappedInstance = wrappedInstance
+  }
+
+  @objc public func getValue(_ opt: String?) -> String? {
+    let result =  wrappedInstance[opt]
+    return result
+  }
+
+}
+
 @objc public class AsyncSubscriptWrapper: NSObject {
   var wrappedInstance: AsyncSubscript
 
@@ -90,20 +104,6 @@
 
 }
 
-@objc public class OptionalSubscriptWrapper: NSObject {
-  var wrappedInstance: OptionalSubscript
-
-  init(_ wrappedInstance: OptionalSubscript) {
-    self.wrappedInstance = wrappedInstance
-  }
-
-  @objc public func getValue(_ opt: IntWrapper?) -> String? {
-    let result =  wrappedInstance[opt?.wrappedInstance]
-    return result
-  }
-
-}
-
 @objc public class ThrowingSubscriptWrapper: NSObject {
   var wrappedInstance: ThrowingSubscript
 
@@ -146,25 +146,16 @@
 
 }
 
-@objc public class NonTrivialArgsSubscriptWrapper: NSObject {
-  var wrappedInstance: NonTrivialArgsSubscript
+@objc public class NonTrivialTypesSubscriptWrapper: NSObject {
+  var wrappedInstance: NonTrivialTypesSubscript
 
-  init(_ wrappedInstance: NonTrivialArgsSubscript) {
+  init(_ wrappedInstance: NonTrivialTypesSubscript) {
     self.wrappedInstance = wrappedInstance
   }
 
-  @objc public func getValue(_ other: SubscriptClassWrapper) -> Int {
+  @objc public func getValue(_ other: SubscriptClassWrapper) -> SubscriptClassWrapper {
     let result =  wrappedInstance[other.wrappedInstance]
-    return result
-  }
-
-}
-
-@objc public class IntWrapper: NSObject {
-  var wrappedInstance: Int
-
-  init(_ wrappedInstance: Int) {
-    self.wrappedInstance = wrappedInstance
+    return SubscriptClassWrapper(result)
   }
 
 }