Change Classifier to accept language name Strings

1 files changed, 14 insertions, 14 deletions

diff --git a/test/test_classifier.rb b/test/test_classifier.rb
index 9a72c9b..33de385 100644
--- a/test/test_classifier.rb
+++ b/test/test_classifier.rb
@@ -25,29 +25,29 @@ class TestClassifier < Test::Unit::TestCase
 
   def test_classify
     classifier = Classifier.new
-    classifier.train Language["Ruby"], fixture("ruby/foo.rb")
-    classifier.train Language["Objective-C"], fixture("objective-c/Foo.h")
-    classifier.train Language["Objective-C"], fixture("objective-c/Foo.m")
+    classifier.train "Ruby", fixture("ruby/foo.rb")
+    classifier.train "Objective-C", fixture("objective-c/Foo.h")
+    classifier.train "Objective-C", fixture("objective-c/Foo.m")
 
     results = classifier.classify(fixture("objective-c/hello.m"))
-    assert_equal Language["Objective-C"], results.first[0]
+    assert_equal "Objective-C", results.first[0]
 
     tokens  = Tokenizer.tokenize(fixture("objective-c/hello.m"))
     results = classifier.classify(tokens)
-    assert_equal Language["Objective-C"], results.first[0]
+    assert_equal "Objective-C", results.first[0]
   end
 
   def test_restricted_classify
     classifier = Classifier.new
-    classifier.train Language["Ruby"], fixture("ruby/foo.rb")
-    classifier.train Language["Objective-C"], fixture("objective-c/Foo.h")
-    classifier.train Language["Objective-C"], fixture("objective-c/Foo.m")
+    classifier.train "Ruby", fixture("ruby/foo.rb")
+    classifier.train "Objective-C", fixture("objective-c/Foo.h")
+    classifier.train "Objective-C", fixture("objective-c/Foo.m")
 
-    results = classifier.classify(fixture("objective-c/hello.m"), [Language["Objective-C"]])
-    assert_equal Language["Objective-C"], results.first[0]
+    results = classifier.classify(fixture("objective-c/hello.m"), ["Objective-C"])
+    assert_equal "Objective-C", results.first[0]
 
-    results = classifier.classify(fixture("objective-c/hello.m"), [Language["Ruby"]])
-    assert_equal Language["Ruby"], results.first[0]
+    results = classifier.classify(fixture("objective-c/hello.m"), ["Ruby"])
+    assert_equal "Ruby", results.first[0]
   end
 
   def test_instance_classify_empty
@@ -72,11 +72,11 @@ class TestClassifier < Test::Unit::TestCase
       next unless sample.language.overrides.any?
 
       extname   = File.extname(sample.path)
-      languages = Language.all.select { |l| l.extensions.include?(extname) }
+      languages = Language.all.select { |l| l.extensions.include?(extname) }.map(&:name)
       next unless languages.length > 1
 
       results = Classifier.instance.classify(sample.data, languages)
-      assert_equal sample.language, results.first[0], "#{sample.path}\n#{results.inspect}"
+      assert_equal sample.language.name, results.first[0], "#{sample.path}\n#{results.inspect}"
     end
   end
 end