diff --git a/lib/membrane_tcp/socket.ex b/lib/membrane_tcp/socket.ex
index 6c53ac0..e2cb95f 100644
--- a/lib/membrane_tcp/socket.ex
+++ b/lib/membrane_tcp/socket.ex
@@ -1,7 +1,9 @@
 defmodule Membrane.TCP.Socket do
-  @moduledoc false
+  @moduledoc """
+  TCP Socket behavior
+  """
 
-  @enforce_keys [:port_no, :ip_address]
+  @enforce_keys [:connection_side, :port_no, :ip_address]
   defstruct [:port_no, :ip_address, :socket_handle, :state, :connection_side, sock_opts: []]
 
   @type t :: %__MODULE__{
@@ -14,7 +16,7 @@ defmodule Membrane.TCP.Socket do
         }
 
   @type socket_pair_config :: %{
-          connection_side: :server | :client | {:client, :inet.address(), :inet.port_number()},
+          connection_side: :server | :client | {:client, :inet.ip_address(), :inet.port_number()},
           local_address: :inet.socket_address(),
           local_port_no: :inet.port_number(),
           local_socket: t() | nil
diff --git a/test/membrane_tcp/integration_test.exs b/test/membrane_tcp/integration_test.exs
index d659155..ab3ef8e 100644
--- a/test/membrane_tcp/integration_test.exs
+++ b/test/membrane_tcp/integration_test.exs
@@ -37,7 +37,12 @@ defmodule Membrane.TCP.IntegrationTest do
       )
 
     assert_pipeline_notified(sender, :sink, {:connection_info, @local_address, @server_port})
-    assert_pipeline_notified(receiver, :source, {:connection_info, @local_address, _ephemeral_port})
+
+    assert_pipeline_notified(
+      receiver,
+      :source,
+      {:connection_info, @local_address, _ephemeral_port}
+    )
 
     assert_end_of_stream(sender, :sink)
 
diff --git a/test/membrane_tcp/socket_test.exs b/test/membrane_tcp/socket_test.exs
index 55e798b..7d690eb 100644
--- a/test/membrane_tcp/socket_test.exs
+++ b/test/membrane_tcp/socket_test.exs
@@ -5,7 +5,7 @@ defmodule Membrane.TCP.SocketTest do
 
   describe "listen" do
     test "with explicit port and address" do
-      sock = %Socket{port_no: 50_666, ip_address: {127, 0, 0, 1}}
+      sock = %Socket{connection_side: :client, port_no: 50_666, ip_address: {127, 0, 0, 1}}
       assert {:ok, new_sock} = Socket.listen(sock)
       assert new_sock.ip_address == sock.ip_address
       assert new_sock.port_no == sock.port_no
@@ -15,7 +15,7 @@ defmodule Membrane.TCP.SocketTest do
     end
 
     test "with port 0 and `:any` IPv6 address" do
-      sock = %Socket{port_no: 0, ip_address: :any, sock_opts: [:inet6]}
+      sock = %Socket{connection_side: :client, port_no: 0, ip_address: :any, sock_opts: [:inet6]}
       assert {:ok, new_sock} = Socket.listen(sock)
       assert new_sock.ip_address == {0, 0, 0, 0, 0, 0, 0, 0}
       assert new_sock.port_no != 0
@@ -25,7 +25,13 @@ defmodule Membrane.TCP.SocketTest do
     end
 
     test "with port 0 and `:loopback` IPv4 address" do
-      sock = %Socket{port_no: 0, ip_address: :loopback, sock_opts: [:inet]}
+      sock = %Socket{
+        connection_side: :client,
+        port_no: 0,
+        ip_address: :loopback,
+        sock_opts: [:inet]
+      }
+
       assert {:ok, new_sock} = Socket.listen(sock)
       assert new_sock.ip_address == {127, 0, 0, 1}
       assert new_sock.port_no != 0
@@ -37,8 +43,13 @@ defmodule Membrane.TCP.SocketTest do
 
   describe "socket pair" do
     test "completes the handshake successfully" do
-      client_socket = %Socket{port_no: 50_667, ip_address: :loopback}
-      server_socket = %Socket{port_no: 50_666, ip_address: {127, 0, 0, 1}}
+      client_socket = %Socket{connection_side: :client, port_no: 50_667, ip_address: :loopback}
+
+      server_socket = %Socket{
+        connection_side: :server,
+        port_no: 50_666,
+        ip_address: {127, 0, 0, 1}
+      }
 
       assert {:ok, listening_server_socket} = Socket.listen(server_socket)
       assert {:ok, client_socket} = Socket.connect(client_socket, server_socket)
diff --git a/test/membrane_tcp/source/pipeline_test.exs b/test/membrane_tcp/source/pipeline_test.exs
index 34e8672..dd6bfcd 100644
--- a/test/membrane_tcp/source/pipeline_test.exs
+++ b/test/membrane_tcp/source/pipeline_test.exs
@@ -28,6 +28,7 @@ defmodule Membrane.TCP.SourcePipelineTest do
     test "created without socket" do
       {:ok, listening_server_socket} =
         %Socket{
+          connection_side: :server,
           ip_address: @local_address,
           port_no: @server_port_no,
           sock_opts: [reuseaddr: true]
@@ -53,13 +54,11 @@ defmodule Membrane.TCP.SourcePipelineTest do
     end
 
     test "created with already connected client socket" do
-      server_socket = %Socket{
-        connection_side: :server,
-        ip_address: @local_address,
-        port_no: @server_port_no
-      }
+      server_socket =
+        %Socket{connection_side: :server, ip_address: @local_address, port_no: @server_port_no}
 
-      client_socket = %Socket{connection_side: :client, ip_address: @local_address, port_no: 0}
+      client_socket =
+        %Socket{connection_side: :client, ip_address: @local_address, port_no: 0}
 
       assert {:ok, listening_server_socket} = Socket.listen(server_socket)
       assert {:ok, client_socket} = Socket.connect(client_socket, server_socket)