From 7bc612efee88a57662f6d56698f8647280470d83 Mon Sep 17 00:00:00 2001
From: Adam Rocska <adam.rocska@adams.solutions>
Date: Thu, 30 Apr 2020 14:43:44 +0200
Subject: [PATCH] ControlByte constructor got changed, because the byte
 sequence is at least 1 byte, at most 5 bytes long.

Signed-off-by: Adam Rocska <adam.rocska@adams.solutions>
---
 Sources/MaxMindDBReader/ControlByte.swift     |  9 ++-
 .../ControlByteTest.swift                     | 67 ++++++++++++-------
 2 files changed, 51 insertions(+), 25 deletions(-)

diff --git a/Sources/MaxMindDBReader/ControlByte.swift b/Sources/MaxMindDBReader/ControlByte.swift
index 0821cfb..d7c0dc2 100644
--- a/Sources/MaxMindDBReader/ControlByte.swift
+++ b/Sources/MaxMindDBReader/ControlByte.swift
@@ -4,14 +4,19 @@ struct ControlByte {
 
   let type: DataType
 
-  init?(firstByte: UInt8, secondByte: UInt8 = 0b0000_0000) {
+  init?(bytes: Data) {
+    if bytes.count == 0 || bytes.count > 5 { return nil }
+
+    let firstByte                 = bytes.first!
     let typeDefinitionOnFirstByte = firstByte &>> 5
+
     guard let type = typeDefinitionOnFirstByte != 0b0000_0000
                      ? DataType(rawValue: typeDefinitionOnFirstByte)
-                     : DataType(rawValue: secondByte + 7)
+                     : DataType(rawValue: bytes[bytes.index(after: bytes.startIndex)] + 7)
       else {
       return nil
     }
     self.type = type
   }
+
 }
diff --git a/Tests/MaxMindDBReaderTests/ControlByteTest.swift b/Tests/MaxMindDBReaderTests/ControlByteTest.swift
index 9d28297..718cb0e 100644
--- a/Tests/MaxMindDBReaderTests/ControlByteTest.swift
+++ b/Tests/MaxMindDBReaderTests/ControlByteTest.swift
@@ -3,31 +3,52 @@ import XCTest
 @testable import MaxMindDBReader
 
 class ControlByteTest: XCTestCase {
-  func testInit() {
-    let dataTypes = (1...255).compactMap({ DataType(rawValue: $0) })
 
-    let nonExtendedTypes = dataTypes.filter({ $0.rawValue <= 7 }).map({ $0.rawValue })
-    let extendedTypes = dataTypes.filter({ $0.rawValue > 7 }).map({ $0.rawValue })
-    precondition(nonExtendedTypes.count > 0, "nonExtendedTypes can't be empty.")
-    precondition(extendedTypes.count > 0, "extendedTypes can't be empty.")
-    for nonExtendedType in nonExtendedTypes {
-      XCTAssertEqual(
-        DataType(rawValue: nonExtendedType),
-        ControlByte(firstByte: nonExtendedType &<< 5)?.type
-      )
-      XCTAssertEqual(
-        DataType(rawValue: nonExtendedType),
-        ControlByte(
-          firstByte: nonExtendedType &<< 5,
-          secondByte: 0b1111_1111
-        )?.type
-      )
+  private static let dataTypes: [DataType] = (1...255).compactMap({ DataType(rawValue: $0) })
+  private static let nonExtendedRawValues: [DataType.RawValue] = dataTypes
+    .filter({ $0.rawValue <= 7 })
+    .map({ $0.rawValue })
+  private static let extendedRawValues: [DataType.RawValue] = dataTypes
+    .filter({ $0.rawValue > 7 })
+    .map({ $0.rawValue })
+
+  override class func setUp() {
+    super.setUp()
+    precondition(nonExtendedRawValues.count > 0, "nonExtendedRawValues can't be empty.")
+    precondition(extendedRawValues.count > 0, "extendedRawValues can't be empty.")
+  }
+
+  func testInit_nilIfEmpty() {
+    XCTAssertNil(ControlByte(bytes: Data()))
+  }
+
+  func testInit_nilIfBiggerThanFive() {
+    XCTAssertNil(ControlByte(bytes: Data([0b0100_0000]) + Data(count: 5)))
+    XCTAssertNil(ControlByte(bytes: Data([0b0100_0000]) + Data(count: 6)))
+    XCTAssertNil(ControlByte(bytes: Data([0b0100_0000]) + Data(count: 7)))
+    XCTAssertNil(ControlByte(bytes: Data([0b0100_0000]) + Data(count: 10)))
+    XCTAssertNil(ControlByte(bytes: Data([0b0100_0000]) + Data(count: 100)))
+    XCTAssertNil(ControlByte(bytes: Data([0b0100_0000]) + Data(count: 10_000)))
+  }
+
+  func testInit_dataTypeIdentification() {
+    for value in ControlByteTest.nonExtendedRawValues {
+      for data in (0...4).map({ Data([value &<< 5]) + Data(count: $0) }) {
+        XCTAssertEqual(
+          DataType(rawValue: value),
+          ControlByte(bytes: data)?.type
+        )
+      }
     }
-    for extendedType in extendedTypes {
-      XCTAssertEqual(
-        DataType(rawValue: extendedType),
-        ControlByte(firstByte: 0b0000_0000, secondByte: extendedType - 7)?.type
-      )
+
+    for value in ControlByteTest.extendedRawValues {
+      for data in (0...3).map({ Data([0b0000_0000, value - 7]) + Data(count: $0) }) {
+        XCTAssertEqual(
+          DataType(rawValue: value),
+          ControlByte(bytes: data)?.type
+        )
+      }
     }
   }
+
 }