diff --git a/Sources/AvatarView.swift b/Sources/AvatarView.swift
index 7f91caab..537cc32a 100644
--- a/Sources/AvatarView.swift
+++ b/Sources/AvatarView.swift
@@ -43,17 +43,13 @@ open class AvatarView: UIView {
     
     func getImageFrom(initals: String, withColor color: UIColor = UIColor.white, fontSize: CGFloat = 14) -> UIImage {
         _ = UIGraphicsBeginImageContext(CGSize(width: 30, height: 30))
-        //// General Declarations
         let context = UIGraphicsGetCurrentContext()!
         
-        //// Color Declarations
-        let white = UIColor.white
-        
         //// Text Drawing
         let textRect = CGRect(x: 5, y: 6, width: 20, height: 20)
         let textStyle = NSMutableParagraphStyle()
         textStyle.alignment = .center
-        let textFontAttributes = [NSFontAttributeName: UIFont.systemFont(ofSize: fontSize), NSForegroundColorAttributeName: white, NSParagraphStyleAttributeName: textStyle]
+        let textFontAttributes = [NSFontAttributeName: UIFont.systemFont(ofSize: fontSize), NSForegroundColorAttributeName: color, NSParagraphStyleAttributeName: textStyle]
         
         let textTextHeight: CGFloat = initals.boundingRect(with: CGSize(width: textRect.width, height: CGFloat.infinity), options: .usesLineFragmentOrigin, attributes: textFontAttributes, context: nil).height
         context.saveGState()
diff --git a/Tests/AvatarViewTests.swift b/Tests/AvatarViewTests.swift
index 1387f5e4..0c089d3a 100644
--- a/Tests/AvatarViewTests.swift
+++ b/Tests/AvatarViewTests.swift
@@ -20,61 +20,51 @@ class AvatarViewTests: XCTestCase {
     }
     
     func testNoParams() {
-        let avatar = AvatarView()
-        XCTAssertEqual(avatar.initals, "?")
-        XCTAssertEqual(avatar.initalsLabel.text, "?")
-        XCTAssertEqual(avatar.layer.cornerRadius, 15.0)
-        XCTAssertNil(avatar.imageView.image)
-        XCTAssertTrue(avatar.imageView.isHidden)
-        XCTAssertEqual(avatar.initalsLabel.textColor, UIColor.white)
+        let avatarView = AvatarView()
+        XCTAssertEqual(avatarView.avatar.initals, "?")
+        XCTAssertEqual(avatarView.layer.cornerRadius, 15.0)
+        XCTAssertEqual(avatarView.backgroundColor, UIColor.gray)
     }
     
     func testWithImage() {
-        let avatar = AvatarView(image: UIImage())
+        let avatarView = AvatarView()
+        let avatar = Avatar(image: UIImage())
+        avatarView.set(avatar: avatar)
         XCTAssertEqual(avatar.initals, "?")
-        XCTAssertEqual(avatar.initalsLabel.text, "?")
-        XCTAssertEqual(avatar.layer.cornerRadius, 15.0)
-        XCTAssertFalse(avatar.imageView.isHidden)
-        XCTAssertEqual(avatar.initalsLabel.textColor, UIColor.white)
-        XCTAssertEqual(avatar.backgroundColor, UIColor.gray)
+        XCTAssertEqual(avatarView.layer.cornerRadius, 15.0)
+        XCTAssertEqual(avatarView.backgroundColor, UIColor.gray)
     }
     
-    func testCustom() {
-        let avatar = AvatarView(size: 50, image: UIImage(), highlightedImage: nil, initals: "lol", cornerRounding: 6)
-        XCTAssertEqual(avatar.initals, "lol")
-        XCTAssertEqual(avatar.initalsLabel.text, "lol")
-        XCTAssertEqual(avatar.layer.cornerRadius, 6.0)
-        XCTAssertFalse(avatar.imageView.isHidden)
-        XCTAssertEqual(avatar.initalsLabel.textColor, UIColor.white)
-        XCTAssertEqual(avatar.backgroundColor, UIColor.gray)
+    func testInitalsOnly() {
+        let avatarView = AvatarView()
+        let avatar = Avatar(initals: "DL")
+        avatarView.set(avatar: avatar)
+        XCTAssertEqual(avatar.initals, "DL")
+        XCTAssertEqual(avatarView.layer.cornerRadius, 15.0)
+        XCTAssertEqual(avatarView.backgroundColor, UIColor.gray)
     }
     
     func testSetBackground() {
-        let avatar = AvatarView(image: UIImage())
-        XCTAssertEqual(avatar.backgroundColor, UIColor.gray)
-        avatar.setBackground(color: UIColor.red)
-        XCTAssertEqual(avatar.backgroundColor, UIColor.red)
+        let avatarView = AvatarView()
+        XCTAssertEqual(avatarView.backgroundColor, UIColor.gray)
+        avatarView.setBackground(color: UIColor.red)
+        XCTAssertEqual(avatarView.backgroundColor, UIColor.red)
     }
     
     func testGetImage() {
         let image = UIImage()
-        let avatar = AvatarView(image: image)
-        XCTAssertEqual(avatar.getImage(), image)
+        let avatar = Avatar(image: image)
+        let avatarView = AvatarView()
+        avatarView.set(avatar: avatar)
+        XCTAssertEqual(avatarView.getImage(), image)
     }
     
     func testRoundedCorners() {
-        let avatar = AvatarView(image: UIImage())
-        XCTAssertEqual(avatar.layer.cornerRadius, 15.0)
-        avatar.setCorner(radius: 2)
-        XCTAssertEqual(avatar.layer.cornerRadius, 2.0)
-    }
-    
-    func testInitalsFont() {
-        let avatar = AvatarView(image: UIImage())
-        XCTAssertEqual(avatar.initalsLabel.textColor, UIColor.white)
-        XCTAssertEqual(avatar.initalsLabel.font.pointSize, 16)
-        avatar.setInitalsFont(size: 20, color: UIColor.blue)
-        XCTAssertEqual(avatar.initalsLabel.textColor, UIColor.blue)
-        XCTAssertEqual(avatar.initalsLabel.font.pointSize, 20)
+        let avatarView = AvatarView()
+        let avatar = Avatar(image: UIImage())
+        avatarView.set(avatar: avatar)
+        XCTAssertEqual(avatarView.layer.cornerRadius, 15.0)
+        avatarView.setCorner(radius: 2)
+        XCTAssertEqual(avatarView.layer.cornerRadius, 2.0)
     }
 }