diff --git a/rtl/cache_memory.sv b/rtl/cache_memory.sv
index aaa9a03..cee5f4d 100644
--- a/rtl/cache_memory.sv
+++ b/rtl/cache_memory.sv
@@ -7,7 +7,7 @@ module cache_memory #(
     parameter int BLOCK_SIZE = 32
 ) (
     input logic clk,
-    rst,
+    input logic rst,
     input logic [WaySize - 1:0] way,
     input logic [SetSize - 1:0] set,
     input logic [TagSize - 1:0] tag,
@@ -29,7 +29,7 @@ module cache_memory #(
     logic valid;
   } cache_line_t;
 
-  typedef cache_line [NUM_SETS - 1:0] cache_way;
+  typedef cache_line_t [NUM_SETS - 1:0] cache_way;
   cache_way [NUM_WAYS - 1:0] ways;
 
   assign read_data = ways[way][set].data;
@@ -39,8 +39,8 @@ module cache_memory #(
       // Reset valid flags
       for (int i = 0; i < NUM_WAYS; i++) begin
         for (int j = 0; j < NUM_SETS; j++) begin
-          ways[i][j].data  <= 'dx;
-          ways[i][j].tag   <= 'dx;
+          ways[i][j].data  <= BLOCK_SIZE'('dx);
+          ways[i][j].tag   <= TagSize'('dx);
           ways[i][j].valid <= 0;
         end
       end
diff --git a/rtl/two_way_lru_cache.sv b/rtl/two_way_lru_cache.sv
index b388af9..4159d2c 100644
--- a/rtl/two_way_lru_cache.sv
+++ b/rtl/two_way_lru_cache.sv
@@ -1,77 +1,78 @@
 `include "timescale.sv"
 
 module two_way_lru_cache #(
-    parameter ADDR_SIZE = 32,
-    parameter NUM_SETS = 16,
-    parameter BLOCK_SIZE = 32
-)(
-    input logic clk, rst,
-    input logic[ADDR_SIZE - 1:0] addr,
+    parameter int ADDR_SIZE  = 32,
+    parameter int NUM_SETS   = 16,
+    parameter int BLOCK_SIZE = 32
+) (
+    input logic clk,
+    input logic rst,
+    input logic [ADDR_SIZE - 1:0] addr,
     input logic write_enable,
-    input logic[BLOCK_SIZE - 1:0] write_data,
-    output logic[BLOCK_SIZE - 1:0] read_data,
+    input logic [BLOCK_SIZE - 1:0] write_data,
+    output logic [BLOCK_SIZE - 1:0] read_data,
     output logic hit
 );
-    localparam NUM_WAYS = 2;
-    localparam NUM_BLOCK_BYTES = BLOCK_SIZE / 4;
-    localparam BYTE_OFFSET_SIZE = $clog2(NUM_BLOCK_BYTES);
-    localparam WAY_SIZE = $clog2(NUM_WAYS);
-    localparam SET_SIZE = $clog2(NUM_SETS);
-    localparam TAG_SIZE = ADDR_SIZE - SET_SIZE - BYTE_OFFSET_SIZE;
-    
-    logic[NUM_WAYS - 1:0] valid_flags;
-    logic[NUM_WAYS - 1:0] hits;
-    
-    logic[WAY_SIZE - 1:0] way;
-    logic[SET_SIZE - 1:0] set;
-    logic[TAG_SIZE - 1:0] tag;
-    
-    CacheMemory #(
-        ADDR_SIZE,
-        NUM_SETS,
-        NUM_WAYS,
-        BLOCK_SIZE
-    ) cache_memory(
-        clk,
-        rst,
-        way,
-        set,
-        tag,
-        write_enable,
-        write_data,
-        read_data,
-        hits,
-        valid_flags
-    );
-    
-    TwoWayLRUCRU #(
-        ADDR_SIZE,
-        NUM_SETS,
-        BLOCK_SIZE
-    ) cache_replace_unit(
-        clk,
-        rst,
-        addr,
-        cru_enable,
-        replace_preferred_way
-    );
-    
-    CacheController #(
-        ADDR_SIZE,
-        NUM_SETS,
-        NUM_WAYS,
-        BLOCK_SIZE
-    ) cache_controller (
-        clk,
-        addr,
-        write_enable,
-        replace_preferred_way,
-        hits,
-        valid_flags,
-        set,
-        tag,
-        way,
-        hit,
-        cru_enable
-    );
+  localparam int NumWays = 2;
+  localparam int NumBlockBytes = BLOCK_SIZE / 4;
+  localparam int ByteOffsetSize = $clog2(NUM_BLOCK_BYTES);
+  localparam int WaySize = $clog2(NUM_WAYS);
+  localparam int SetSize = $clog2(NUM_SETS);
+  localparam int TagSize = ADDR_SIZE - SET_SIZE - BYTE_OFFSET_SIZE;
+
+  logic [NUM_WAYS - 1:0] valid_flags;
+  logic [NUM_WAYS - 1:0] hits;
+
+  logic [WAY_SIZE - 1:0] way;
+  logic [SET_SIZE - 1:0] set;
+  logic [TAG_SIZE - 1:0] tag;
+
+  cache_memory #(
+      .ADDR_SIZE (ADDR_SIZE),
+      .NUM_SETS  (NUM_SETS),
+      .NUM_WAYS  (NumWays),
+      .BLOCK_SIZE(BLOCK_SIZE)
+  ) cache_memory (
+      .clk(clk),
+      .rst(rst),
+      .way(way),
+      .set(set),
+      .tag(tag),
+      .write_enable(write_enable),
+      .write_data(write_data),
+      .read_data(read_data),
+      .hits(hits),
+      .valid_flags(valid_flags)
+  );
+
+  two_way_lru_cru #(
+      .ADDR_SIZE (ADDR_SIZE),
+      .NUM_SETS  (NUM_SETS),
+      .BLOCK_SIZE
+  ) cache_replace_unit (
+      .clk(clk),
+      .rst(rst),
+      .addr(addr),
+      .replace(cru_enable),
+      .preferred(replace_preferred_way)
+  );
+
+  cache_controller #(
+      .ADDR_SIZE (ADDR_SIZE),
+      .NUM_SETS  (NUM_SETS),
+      .NUM_WAYS  (NUM_WAYS),
+      .BLOCK_SIZE(BLOCK_SIZE)
+  ) cache_controller (
+      .clk(clk),
+      .addr(addr),
+      .write_enable(write_enable),
+      .replace_way(replace_preferred_way),
+      .hits(hits),
+      .valid_flags(valid_flags),
+      .set(set),
+      .tag(tag),
+      .way(way),
+      .hit(hit),
+      .cru_enable(cru_enable)
+  );
 endmodule
diff --git a/rtl/two_way_lru_cru.sv b/rtl/two_way_lru_cru.sv
index 79a3d69..b66d455 100644
--- a/rtl/two_way_lru_cru.sv
+++ b/rtl/two_way_lru_cru.sv
@@ -1,37 +1,37 @@
 `include "timescale.sv"
 
 module two_way_lru_cru #(
-    parameter ADDR_SIZE = 32,
-    parameter NUM_SETS = 16,
-    parameter BLOCK_SIZE = 32
+    parameter int ADDR_SIZE  = 32,
+    parameter int NUM_SETS   = 16,
+    parameter int BLOCK_SIZE = 32
 ) (
-    input logic clk, rst,
-    input logic[ADDR_SIZE - 1:0] addr,
+    input logic clk,
+    input logic rst,
+    input logic [ADDR_SIZE - 1:0] addr,
     input logic replace,
     output logic preferred
 );
-    localparam NUM_BLOCK_BYTES = BLOCK_SIZE / 4;
-    localparam BYTE_OFFSET_SIZE = $clog2(NUM_BLOCK_BYTES);
-    localparam SET_SIZE = $clog2(NUM_SETS);
-    localparam TAG_SIZE = ADDR_SIZE - SET_SIZE - BYTE_OFFSET_SIZE;
-    
-    typedef struct packed {
-        logic[BYTE_OFFSET_SIZE - 1:0] byte_offset;
-        logic[SET_SIZE - 1:0] set;
-        logic[TAG_SIZE - 1:0] tag;
-    } cache_addr;
-    
-    cache_addr packed_addr;
-    assign packed_addr = cache_addr'(addr);
-    
-    logic[NUM_SETS - 1:0] lru;
-    
-    assign preferred = lru[packed_addr.set];
-    always_ff @(posedge clk) begin
-        if (rst)
-            lru[packed_addr.set] <= 0;
-        else if (replace) begin
-            lru[packed_addr.set] <= !lru[packed_addr.set];
-        end
+  localparam int NumBlocksBytes = BLOCK_SIZE / 4;
+  localparam int ByteOffsetSize = $clog2(NumBlocksBytes);
+  localparam int SetSize = $clog2(NumSets);
+  localparam int TagSize = ADDR_SIZE - SetSize - ByteOffSetSize;
+
+  typedef struct packed {
+    logic [ByteOffsetSize - 1:0] byte_offset;
+    logic [SetSize - 1:0] set;
+    logic [TagSize - 1:0] tag;
+  } cache_addr_t;
+
+  cache_addr_t packed_addr;
+  assign packed_addr = cache_addr_t'(addr);
+
+  logic [NUM_SETS - 1:0] lru;
+
+  assign preferred = lru[packed_addr.set];
+  always_ff @(posedge clk) begin
+    if (rst) lru[packed_addr.set] <= 0;
+    else if (replace) begin
+      lru[packed_addr.set] <= !lru[packed_addr.set];
     end
+  end
 endmodule
diff --git a/test/CMakeLists.txt b/test/CMakeLists.txt
index 97cfc15..c27af7b 100644
--- a/test/CMakeLists.txt
+++ b/test/CMakeLists.txt
@@ -96,3 +96,10 @@ rvscc_add_test(
           ${PROJECT_SOURCE_DIR}/rtl/hazard_unit.sv  
   		  ${PROJECT_SOURCE_DIR}/test/test_five_stage_pipeline_core.sv
 )
+
+rvscc_add_test(
+	NAME cache-memory
+	TOP test_cache_memory
+	SOURCES ${PROJECT_SOURCE_DIR}/rtl/cache_memory.sv
+	        ${PROJECT_SOURCE_DIR}/test/test_cache_memory.sv
+)
diff --git a/test/test_cache_memory.sv b/test/test_cache_memory.sv
index 3be7d26..a52358a 100644
--- a/test/test_cache_memory.sv
+++ b/test/test_cache_memory.sv
@@ -1,29 +1,64 @@
-`timescale 1ns / 1ps
+`include "timescale.sv"
 
 module test_cache_memory ();
+  logic clk;
+  logic rst;
 
-  logic [31:0] addr, write_data, read_data;
-  logic clk, rst, write_enable;
-
-  CacheMemory cache_memory (
+  logic [dut.WaySize-1:0] way;
+  logic [dut.SetSize-1:0] set;
+  logic [dut.TagSize-1:0] tag;
+  logic write_enable;
+  logic [31:0] write_data;
+  logic [31:0] read_data;
+  logic [1:0] hits;
+  logic [1:0] valid_flags;
+  cache_memory #(
+      .ADDR_SIZE (32),
+      .NUM_SETS  (4),
+      .NUM_WAYS  (2),
+      .BLOCK_SIZE(32)
+  ) dut (
       .clk(clk),
       .rst(rst),
-      write_enable,
-      write_data,
-      read_data
+      .way(way),
+      .set(set),
+      .tag(tag),
+      .write_enable(write_enable),
+      .write_data(write_data),
+      .read_data(read_data),
+      .hits(hits),
+      .valid_flags(valid_flags)
   );
 
-  always #5 clk = ~clk;
+  localparam int ClockCycle = 2;
+  always #ClockCycle clk = !clk;
+
+  logic [31:0] write_value;
 
   initial begin
+    $dumpfile("cache.vcd");
+    $dumpvars;
     clk = 0;
     rst = 1;
-    write_enable = 0;
-    #25 rst = 0;
-    addr = 'd7;
+    #ClockCycle;
+    rst = 0;
+
+    way = 0;
+    set = 0;
+    tag = 27'($urandom);
     write_enable = 1;
-    write_data = 'd10;
-    #25 write_enable = 0;
+    write_value = $urandom;
+    write_data = write_value;
+    #ClockCycle;
+    write_enable = 0;
+    tag += 1;
+    assert (valid_flags == 'b00)
+    else $error("Valid flags does not match");
+    #1;
+    tag -= 1;
+    assert (valid_flags == 'b01)
+    else $error("Valid flags does not match");
+    $finish;
   end
 
 endmodule