Merge pull request #5 from daehn/fix/add-when-non-existent

Add non-existent elements when calling setCount(_, for:)
daehn · Sep 15, 2017 · 397aa51 · 397aa51
2 parents 35405f0 + a28f853
commit 397aa51
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 16 deletions.
diff --git a/CountedSet/CountedSet.swift b/CountedSet/CountedSet.swift
@@ -55,12 +55,15 @@ public struct CountedSet<Element : Hashable> : ExpressibleByArrayLiteral {
         return count(for: member)
     }
 
-    @discardableResult public mutating func setCount(_ count: Int, for element: Element) -> Bool {
+    public mutating func setCount(_ count: Int, for element: Element) {
         precondition(count >= 0, "Count has to be positive")
-        guard contains(element) else { return false }
+        if count > 0 && !contains(element) {
+            backing.insert(element)
+        }
         countByElement[element] = count
-        if count <= 0 { backing.remove(element) }
-        return true
+        if count <= 0 {
+            backing.remove(element)
+        }
     }
 
     public func mostFrequent() -> ElementWithCount? {

diff --git a/CountedSetTests/CountedSetTests.swift b/CountedSetTests/CountedSetTests.swift
@@ -247,34 +247,30 @@ class CountedSetTests: XCTestCase {
         XCTAssertEqual(sut["Baz"], 0)
     }
 
-    func testThatItReturnsFalseWhenSettingCountForNonContainedElement() {
-        // given
-        sut = CountedSet(arrayLiteral: "Foo", "Bar")
-        // then
-        XCTAssertFalse(sut.setCount(2, for: "Baz"))
-    }
-
-    func testThatItReturnsTrueWhenSettingCountForContainedElement() {
+    func testThatItAddsANonExistentElementWhenSettingCount() {
         // given
         sut = CountedSet(arrayLiteral: "Foo", "Bar")
+        // when
+        sut.setCount(2, for: "Baz")
         // then
-        XCTAssertTrue(sut.setCount(2, for: "Foo"))
+        XCTAssert(sut.contains("Baz"))
+        XCTAssertEqual(sut.count(for: "Baz"), 2)
     }
 
     func testThatItSetsTheNewCountWhenSettingCountForContainedElement() {
         // given
         sut = CountedSet(arrayLiteral: "Foo", "Bar")
         // when
-        XCTAssertTrue(sut.setCount(2, for: "Foo"))
+        sut.setCount(3, for: "Foo")
         // then
-        XCTAssertEqual(sut.count(for: "Foo"), 2)
+        XCTAssertEqual(sut.count(for: "Foo"), 3)
     }
 
     func testThatItRemovesElementFrombackingStoreWhenSettingCountForContainedElementToZero() {
         // given
         sut = CountedSet(arrayLiteral: "Foo", "Bar")
         // when
-        XCTAssertTrue(sut.setCount(0, for: "Foo"))
+        sut.setCount(0, for: "Foo")
         // then
         XCTAssertFalse(sut.contains("Foo"))
         XCTAssertEqual(sut.count(for: "Foo"), 0)