assertions/version.rb in yes_ship_it-0.0.2 vs assertions/version.rb in yes_ship_it-0.0.3

@@ -1,33 +1,44 @@
 module YSI
   class Version < Assertion
-    attr_accessor :error
+    parameter :version_file, "lib/version.rb"
 
-    def initialize(engine)
-      @engine = engine
-    end
-
     def display_name
       "version number"
     end
 
     def check
-      version_file = "lib/version.rb"
       if !File.exist?(version_file)
         @error = "Expected version in #{version_file}"
         return nil
       end
 
-      File.read(version_file).each_line do |line|
-        if line =~ /VERSION = "(.*)"/
-          @engine.version = $1
-          return @engine.version
-        end
+      version = parse_version(version_file)
+      if version
+        @engine.version = version
+        return @engine.version
       end
       @error = "Couldn't find version in #{version_file}"
       nil
     end
 
     def assert(dry_run: false)
+    end
+
+    def parse_version(file_name)
+      if file_name =~ /\.rb$/
+        File.read(file_name).each_line do |line|
+          if line =~ /VERSION = "(.*)"/
+            return $1
+          end
+        end
+      elsif file_name =~ /\.go$/
+        File.read(file_name).each_line do |line|
+          if line =~ /Version = "(.*)"/
+            return $1
+          end
+        end
+      end
+      nil
     end
   end
 end