diff --git a/agent.go b/agent.go
index 251c7ab..af1af63 100644
--- a/agent.go
+++ b/agent.go
@@ -25,6 +25,8 @@ var doSelfUpdate = selfUpdate
 
 var errRestartRequested = fmt.Errorf("restart requested")
 var joinTimeout = 10 * time.Second
+var heartbeatInterval = 60 * time.Second
+var channelHeartbeatInterval = 25 * time.Second
 
 const maxJobPayloadBytes = 4 << 20 // 4 MB — well above any legitimate job list
 
@@ -224,9 +226,9 @@ func runSession(ctx context.Context, baseURL, token string) error {
 		}
 	}()
 
-	heartbeatTicker := time.NewTicker(60 * time.Second)
+	heartbeatTicker := time.NewTicker(heartbeatInterval)
 	defer heartbeatTicker.Stop()
-	channelHeartbeatTicker := time.NewTicker(25 * time.Second)
+	channelHeartbeatTicker := time.NewTicker(channelHeartbeatInterval)
 	defer channelHeartbeatTicker.Stop()
 	flushTicker := time.NewTicker(100 * time.Millisecond)
 	defer flushTicker.Stop()
diff --git a/agent_test.go b/agent_test.go
index de9e1f3..843bc35 100644
--- a/agent_test.go
+++ b/agent_test.go
@@ -3,8 +3,10 @@ package main
 import (
 	"context"
 	"encoding/base64"
+	"encoding/binary"
 	"encoding/json"
 	"fmt"
+	"io"
 	"net"
 	"strings"
 	"sync"
@@ -284,6 +286,59 @@ func TestHandleMessage(t *testing.T) {
 		// Should just log and not panic
 	})
 
+	t.Run("check_jobs valid", func(t *testing.T) {
+		checkList := &pb.CheckList{Checks: []*pb.Check{
+			{Id: "c1", CheckType: "tcp", TimeoutMs: 1000,
+				Config: &pb.Check_Tcp{Tcp: &pb.TcpCheckConfig{Host: "127.0.0.1", Port: 1}}},
+		}}
+		bin, _ := proto.Marshal(checkList)
+		payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString(bin)})
+
+		snmpCh := make(chan *pb.SnmpResult, 1)
+		mtCh := make(chan *pb.MikrotikResult, 1)
+		credCh := make(chan *pb.CredentialTestResult, 1)
+		monCh := make(chan *pb.MonitoringCheck, 1)
+		checkCh := make(chan *pb.CheckResult, 1)
+		handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh)
+		select {
+		case <-checkCh:
+		case <-time.After(5 * time.Second):
+			t.Error("timed out waiting for check result")
+		}
+	})
+
+	t.Run("check_jobs invalid json", func(t *testing.T) {
+		snmpCh := make(chan *pb.SnmpResult, 1)
+		mtCh := make(chan *pb.MikrotikResult, 1)
+		credCh := make(chan *pb.CredentialTestResult, 1)
+		monCh := make(chan *pb.MonitoringCheck, 1)
+		checkCh := make(chan *pb.CheckResult, 1)
+		handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: json.RawMessage(`not json`)}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh)
+		// Should log error but not panic
+	})
+
+	t.Run("check_jobs invalid base64", func(t *testing.T) {
+		snmpCh := make(chan *pb.SnmpResult, 1)
+		mtCh := make(chan *pb.MikrotikResult, 1)
+		credCh := make(chan *pb.CredentialTestResult, 1)
+		monCh := make(chan *pb.MonitoringCheck, 1)
+		checkCh := make(chan *pb.CheckResult, 1)
+		payload, _ := json.Marshal(map[string]string{"binary": "not-base64!!!"})
+		handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh)
+		// Should log error but not panic
+	})
+
+	t.Run("check_jobs invalid protobuf", func(t *testing.T) {
+		snmpCh := make(chan *pb.SnmpResult, 1)
+		mtCh := make(chan *pb.MikrotikResult, 1)
+		credCh := make(chan *pb.CredentialTestResult, 1)
+		monCh := make(chan *pb.MonitoringCheck, 1)
+		checkCh := make(chan *pb.CheckResult, 1)
+		payload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString([]byte{0xFF, 0xFF, 0xFF})})
+		handleMessage(context.Background(), channelMsg{Event: "check_jobs", Payload: payload}, testPools(t), snmpCh, mtCh, credCh, monCh, checkCh)
+		// Should log error but not panic
+	})
+
 	t.Run("discovery_job event", func(t *testing.T) {
 		origDial := snmpDial
 		defer func() { snmpDial = origDial }()
@@ -698,3 +753,765 @@ func makeTextFrame(payload []byte) []byte {
 	frame = append(frame, payload...)
 	return frame
 }
+
+func TestDispatchJobCancelledContext(t *testing.T) {
+	// submit returns false when context is cancelled, triggering "pool full" log
+	p := testPools(t)
+	snmpCh := make(chan *pb.SnmpResult, 1)
+	mtCh := make(chan *pb.MikrotikResult, 1)
+	credCh := make(chan *pb.CredentialTestResult, 1)
+	monCh := make(chan *pb.MonitoringCheck, 1)
+	checkCh := make(chan *pb.CheckResult, 1)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+
+	// All dispatch types should handle ctx cancellation gracefully
+	dispatchJob(ctx, &pb.AgentJob{
+		JobId: "s1", JobType: pb.JobType_POLL, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
+	}, p, snmpCh, mtCh, credCh, monCh, checkCh)
+
+	dispatchJob(ctx, &pb.AgentJob{
+		JobId: "m1", JobType: pb.JobType_MIKROTIK, MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1"},
+	}, p, snmpCh, mtCh, credCh, monCh, checkCh)
+
+	dispatchJob(ctx, &pb.AgentJob{
+		JobId: "tc1", JobType: pb.JobType_TEST_CREDENTIALS, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
+	}, p, snmpCh, mtCh, credCh, monCh, checkCh)
+
+	dispatchJob(ctx, &pb.AgentJob{
+		JobId: "p1", JobType: pb.JobType_PING, SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
+	}, p, snmpCh, mtCh, credCh, monCh, checkCh)
+}
+
+// fakeWSServer is a test helper that sets up a WebSocket server for runSession tests.
+type fakeWSServer struct {
+	ln   net.Listener
+	conn net.Conn
+}
+
+func newFakeWSServer(t *testing.T) *fakeWSServer {
+	t.Helper()
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _ = ln.Close() })
+	return &fakeWSServer{ln: ln}
+}
+
+func (s *fakeWSServer) addr() string { return s.ln.Addr().String() }
+
+// acceptAndJoin accepts one WS connection and responds with a successful join.
+func (s *fakeWSServer) acceptAndJoin(t *testing.T) {
+	t.Helper()
+	conn, err := s.ln.Accept()
+	if err != nil {
+		t.Logf("accept: %v", err)
+		return
+	}
+	s.conn = conn
+
+	// Read HTTP upgrade
+	buf := make([]byte, 4096)
+	n, _ := conn.Read(buf)
+	key := extractWSKey(string(buf[:n]))
+	accept := computeAcceptKey(key)
+	resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+	_, _ = conn.Write([]byte(resp))
+
+	// Read join frame
+	frameBuf := make([]byte, 4096)
+	_, _ = conn.Read(frameBuf)
+
+	// Send join OK
+	reply, _ := json.Marshal(channelMsg{
+		Topic:   "agent:agent-0",
+		Event:   "phx_reply",
+		Payload: json.RawMessage(`{"status":"ok"}`),
+		Ref:     strPtr("1"),
+	})
+	_, _ = conn.Write(makeTextFrame(reply))
+}
+
+// sendEvent sends a channel message to the connected client.
+func (s *fakeWSServer) sendEvent(event string, payload json.RawMessage) {
+	msg, _ := json.Marshal(channelMsg{
+		Topic:   "agent:agent-0",
+		Event:   event,
+		Payload: payload,
+	})
+	_, _ = s.conn.Write(makeTextFrame(msg))
+}
+
+// close shuts down the server connection.
+func (s *fakeWSServer) close() {
+	if s.conn != nil {
+		_ = s.conn.Close()
+	}
+}
+
+func TestRunSessionCtxCancel(t *testing.T) {
+	srv := newFakeWSServer(t)
+
+	go srv.acceptAndJoin(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	done := make(chan error, 1)
+	go func() {
+		done <- runSession(ctx, "ws://"+srv.addr(), "token")
+	}()
+
+	// Give the session time to enter the main loop
+	time.Sleep(200 * time.Millisecond)
+	cancel()
+
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Errorf("expected nil error on ctx cancel, got: %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Error("runSession did not exit after ctx cancel")
+	}
+	srv.close()
+}
+
+func TestRunSessionReadError(t *testing.T) {
+	srv := newFakeWSServer(t)
+
+	go func() {
+		srv.acceptAndJoin(t)
+		// Small delay, then close to trigger read error
+		time.Sleep(200 * time.Millisecond)
+		srv.close()
+	}()
+
+	err := runSession(context.Background(), "ws://"+srv.addr(), "token")
+	if err == nil {
+		t.Error("expected read error")
+	}
+	if !strings.Contains(err.Error(), "read:") {
+		t.Errorf("expected 'read:' in error, got: %v", err)
+	}
+}
+
+func TestRunSessionInvalidMessage(t *testing.T) {
+	srv := newFakeWSServer(t)
+
+	go func() {
+		srv.acceptAndJoin(t)
+		// Send invalid JSON — should be logged but not crash
+		time.Sleep(100 * time.Millisecond)
+		_, _ = srv.conn.Write(makeTextFrame([]byte("not json")))
+		// Then close to end session
+		time.Sleep(100 * time.Millisecond)
+		srv.close()
+	}()
+
+	err := runSession(context.Background(), "ws://"+srv.addr(), "token")
+	// Should end with read error from close, not crash
+	if err == nil {
+		t.Error("expected error")
+	}
+}
+
+func TestRunSessionConnectError(t *testing.T) {
+	err := runSession(context.Background(), "ws://127.0.0.1:1", "token")
+	if err == nil {
+		t.Error("expected connect error")
+	}
+	if !strings.Contains(err.Error(), "connect:") {
+		t.Errorf("expected 'connect:' in error, got: %v", err)
+	}
+}
+
+func TestRunSessionJoinWriteError(t *testing.T) {
+	// Server accepts WS upgrade but closes immediately after
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		// Read HTTP upgrade
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		key := extractWSKey(string(buf[:n]))
+		accept := computeAcceptKey(key)
+		resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+		// Close immediately so join write may fail
+		_ = conn.Close()
+	}()
+
+	err = runSession(context.Background(), "ws://"+ln.Addr().String(), "token")
+	if err == nil {
+		t.Error("expected error")
+	}
+}
+
+func TestRunSessionJoinUnmarshalError(t *testing.T) {
+	// Server sends binary garbage as join reply
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		key := extractWSKey(string(buf[:n]))
+		accept := computeAcceptKey(key)
+		resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+		frameBuf := make([]byte, 4096)
+		_, _ = conn.Read(frameBuf)
+		// Send invalid JSON as reply
+		_, _ = conn.Write(makeTextFrame([]byte("{invalid json")))
+		time.Sleep(time.Second)
+	}()
+
+	err = runSession(context.Background(), "ws://"+ln.Addr().String(), "token")
+	if err == nil {
+		t.Error("expected unmarshal error")
+	}
+	if !strings.Contains(err.Error(), "join reply unmarshal") {
+		t.Errorf("expected 'join reply unmarshal' in error, got: %v", err)
+	}
+}
+
+func TestRunSessionReadErrorDuringJoin(t *testing.T) {
+	// Server sends upgrade then closes before sending join reply
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		key := extractWSKey(string(buf[:n]))
+		accept := computeAcceptKey(key)
+		resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+		frameBuf := make([]byte, 4096)
+		_, _ = conn.Read(frameBuf)
+		// Close without sending reply
+		_ = conn.Close()
+	}()
+
+	err = runSession(context.Background(), "ws://"+ln.Addr().String(), "token")
+	if err == nil {
+		t.Error("expected read during join error")
+	}
+	if !strings.Contains(err.Error(), "read during join") {
+		t.Errorf("expected 'read during join' in error, got: %v", err)
+	}
+}
+
+func TestRunAgentReconnectOnError(t *testing.T) {
+	// Server that fails first connection then succeeds, then sends restart
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	connCount := 0
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			connCount++
+			buf := make([]byte, 4096)
+			n, _ := conn.Read(buf)
+			key := extractWSKey(string(buf[:n]))
+			accept := computeAcceptKey(key)
+
+			if connCount == 1 {
+				// First connection: upgrade then close immediately
+				resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+				_, _ = conn.Write([]byte(resp))
+				frameBuf := make([]byte, 4096)
+				_, _ = conn.Read(frameBuf)
+				_ = conn.Close()
+			} else {
+				// Second connection: proper session with restart
+				resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+				_, _ = conn.Write([]byte(resp))
+				frameBuf := make([]byte, 4096)
+				_, _ = conn.Read(frameBuf)
+				reply, _ := json.Marshal(channelMsg{
+					Topic:   "agent:agent-0",
+					Event:   "phx_reply",
+					Payload: json.RawMessage(`{"status":"ok"}`),
+					Ref:     strPtr("1"),
+				})
+				_, _ = conn.Write(makeTextFrame(reply))
+				time.Sleep(50 * time.Millisecond)
+				restart, _ := json.Marshal(channelMsg{
+					Topic:   "agent:agent-0",
+					Event:   "restart",
+					Payload: json.RawMessage(`{}`),
+				})
+				_, _ = conn.Write(makeTextFrame(restart))
+				time.Sleep(time.Second)
+				_ = conn.Close()
+			}
+		}
+	}()
+
+	origExit := osExit
+	defer func() { osExit = origExit }()
+	exitCalled := make(chan int, 1)
+	osExit = func(code int) { exitCalled <- code }
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	done := make(chan struct{})
+	go func() {
+		runAgent(ctx, "ws://"+ln.Addr().String(), "token")
+		close(done)
+	}()
+
+	select {
+	case code := <-exitCalled:
+		if code != 0 {
+			t.Errorf("expected exit code 0, got %d", code)
+		}
+	case <-time.After(10 * time.Second):
+		t.Error("runAgent did not reconnect and restart")
+	}
+}
+
+func TestRunAgentContextCancellation(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel() // cancel immediately
+
+	done := make(chan struct{})
+	go func() {
+		runAgent(ctx, "ws://127.0.0.1:1", "token")
+		close(done)
+	}()
+
+	select {
+	case <-done:
+		// runAgent returned as expected
+	case <-time.After(5 * time.Second):
+		t.Error("runAgent did not return after context cancellation")
+	}
+}
+
+func TestRunAgentRestart(t *testing.T) {
+	origExit := osExit
+	defer func() { osExit = origExit }()
+
+	exitCalled := make(chan int, 1)
+	osExit = func(code int) {
+		exitCalled <- code
+	}
+
+	// Start a fake WebSocket server that accepts the join and sends restart
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+
+		buf := make([]byte, 4096)
+		n, _ := conn.Read(buf)
+		reqStr := string(buf[:n])
+		key := extractWSKey(reqStr)
+		accept := computeAcceptKey(key)
+		resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\nSec-WebSocket-Accept: " + accept + "\r\n\r\n"
+		_, _ = conn.Write([]byte(resp))
+
+		// Read join frame
+		frameBuf := make([]byte, 4096)
+		_, _ = conn.Read(frameBuf)
+
+		// Send join OK reply
+		reply, _ := json.Marshal(channelMsg{
+			Topic:   "agent:agent-0",
+			Event:   "phx_reply",
+			Payload: json.RawMessage(`{"status":"ok"}`),
+			Ref:     strPtr("1"),
+		})
+		_, _ = conn.Write(makeTextFrame(reply))
+
+		// Small delay then send restart event
+		time.Sleep(50 * time.Millisecond)
+		restart, _ := json.Marshal(channelMsg{
+			Topic:   "agent:agent-0",
+			Event:   "restart",
+			Payload: json.RawMessage(`{}`),
+		})
+		_, _ = conn.Write(makeTextFrame(restart))
+
+		// Keep connection open briefly
+		time.Sleep(2 * time.Second)
+	}()
+
+	addr := ln.Addr().String()
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	done := make(chan struct{})
+	go func() {
+		runAgent(ctx, "ws://"+addr, "test-token")
+		close(done)
+	}()
+
+	select {
+	case code := <-exitCalled:
+		if code != 0 {
+			t.Errorf("expected exit code 0, got %d", code)
+		}
+	case <-time.After(5 * time.Second):
+		t.Error("osExit was not called after restart")
+	}
+}
+
+// readMaskedFrame reads a single masked WebSocket frame from the server side.
+func readMaskedFrame(conn net.Conn) ([]byte, error) {
+	var header [2]byte
+	if _, err := io.ReadFull(conn, header[:]); err != nil {
+		return nil, err
+	}
+	masked := header[1]&0x80 != 0
+	length := uint64(header[1] & 0x7F)
+	switch length {
+	case 126:
+		var ext [2]byte
+		if _, err := io.ReadFull(conn, ext[:]); err != nil {
+			return nil, err
+		}
+		length = uint64(binary.BigEndian.Uint16(ext[:]))
+	case 127:
+		var ext [8]byte
+		if _, err := io.ReadFull(conn, ext[:]); err != nil {
+			return nil, err
+		}
+		length = binary.BigEndian.Uint64(ext[:])
+	}
+	var maskKey [4]byte
+	if masked {
+		if _, err := io.ReadFull(conn, maskKey[:]); err != nil {
+			return nil, err
+		}
+	}
+	payload := make([]byte, length)
+	if length > 0 {
+		if _, err := io.ReadFull(conn, payload); err != nil {
+			return nil, err
+		}
+	}
+	if masked {
+		for i := range payload {
+			payload[i] ^= maskKey[i%4]
+		}
+	}
+	return payload, nil
+}
+
+// drainFrames reads and discards client frames until stop channel is closed or error.
+func drainFrames(conn net.Conn, stop <-chan struct{}) {
+	for {
+		select {
+		case <-stop:
+			return
+		default:
+		}
+		_ = conn.SetReadDeadline(time.Now().Add(200 * time.Millisecond))
+		if _, err := readMaskedFrame(conn); err != nil {
+			select {
+			case <-stop:
+				return
+			default:
+				continue // timeout, retry
+			}
+		}
+	}
+}
+
+func TestRunSessionProcessesJobResults(t *testing.T) {
+	// Mock external dependencies for fast execution
+	origSnmpDial := snmpDial
+	origMtDial := mikrotikDial
+	origPing := doPing
+	defer func() {
+		snmpDial = origSnmpDial
+		mikrotikDial = origMtDial
+		doPing = origPing
+	}()
+
+	snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) {
+		return &mockSnmpQuerier{
+			getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) {
+				return &gosnmp.SnmpPacket{}, nil
+			},
+		}, func() {}, nil
+	}
+	mikrotikDial = func(ip string, port uint32, username, password string, useSSL bool) (*mikrotikClient, error) {
+		return nil, fmt.Errorf("unreachable")
+	}
+	doPing = func(ip string, timeoutMs int) (float64, error) {
+		return 1.5, nil
+	}
+
+	srv := newFakeWSServer(t)
+	serverDone := make(chan struct{})
+	go func() {
+		defer close(serverDone)
+		srv.acceptAndJoin(t)
+
+		// Drain client frames so writes don't block
+		stopDrain := make(chan struct{})
+		go drainFrames(srv.conn, stopDrain)
+
+		time.Sleep(100 * time.Millisecond)
+
+		// SNMP job → snmpResultCh → snmpBatch → flushSnmpBatch → sendBinaryResult
+		srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{
+			JobId: "s1", JobType: pb.JobType_POLL,
+			SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
+		}))
+
+		// Mikrotik job → mikrotikResultCh → sendBinaryResult("mikrotik_result")
+		srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{
+			JobId: "m1", JobType: pb.JobType_MIKROTIK,
+			MikrotikDevice: &pb.MikrotikDevice{Ip: "10.0.0.1", Port: 8728},
+		}))
+
+		// Ping job → monitoringCheckCh → sendBinaryResult("monitoring_check")
+		srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{
+			JobId: "p1", JobType: pb.JobType_PING, DeviceId: "dev1",
+			SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
+		}))
+
+		// Credential test → credTestResultCh → sendBinaryResult("credential_test_result")
+		srv.sendEvent("jobs", makeJobPayload(&pb.AgentJob{
+			JobId: "ct1", JobType: pb.JobType_TEST_CREDENTIALS,
+			SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1"},
+		}))
+
+		// Check job → checkResultCh → sendBinaryResult("check_result")
+		checkList := &pb.CheckList{Checks: []*pb.Check{
+			{Id: "c1", CheckType: "tcp", TimeoutMs: 500,
+				Config: &pb.Check_Tcp{Tcp: &pb.TcpCheckConfig{Host: "127.0.0.1", Port: 1}}},
+		}}
+		bin, _ := proto.Marshal(checkList)
+		checkPayload, _ := json.Marshal(map[string]string{"binary": base64.StdEncoding.EncodeToString(bin)})
+		srv.sendEvent("check_jobs", checkPayload)
+
+		// Wait for results to flow through all channels
+		time.Sleep(2 * time.Second)
+		close(stopDrain)
+		srv.close()
+	}()
+
+	err := runSession(context.Background(), "ws://"+srv.addr(), "token")
+	<-serverDone
+	if err == nil {
+		t.Error("expected error after server close")
+	}
+}
+
+func TestRunSessionSnmpBatchThreshold(t *testing.T) {
+	origSnmpDial := snmpDial
+	defer func() { snmpDial = origSnmpDial }()
+
+	snmpDial = func(dev *pb.SnmpDevice) (snmpQuerier, func(), error) {
+		return &mockSnmpQuerier{
+			getFunc: func(oids []string) (*gosnmp.SnmpPacket, error) {
+				return &gosnmp.SnmpPacket{}, nil
+			},
+		}, func() {}, nil
+	}
+
+	srv := newFakeWSServer(t)
+	serverDone := make(chan struct{})
+	go func() {
+		defer close(serverDone)
+		srv.acceptAndJoin(t)
+
+		stopDrain := make(chan struct{})
+		go drainFrames(srv.conn, stopDrain)
+
+		time.Sleep(100 * time.Millisecond)
+
+		// Send 55 SNMP jobs to trigger batch threshold (>=50)
+		jobs := make([]*pb.AgentJob, 55)
+		for i := range jobs {
+			jobs[i] = &pb.AgentJob{
+				JobId:      fmt.Sprintf("s%d", i),
+				JobType:    pb.JobType_POLL,
+				SnmpDevice: &pb.SnmpDevice{Ip: "10.0.0.1", Port: 161},
+			}
+		}
+		srv.sendEvent("jobs", makeJobPayload(jobs...))
+
+		time.Sleep(3 * time.Second)
+		close(stopDrain)
+		srv.close()
+	}()
+
+	err := runSession(context.Background(), "ws://"+srv.addr(), "token")
+	<-serverDone
+	if err == nil {
+		t.Error("expected error after server close")
+	}
+}
+
+func TestExecuteCheckPoolFull(t *testing.T) {
+	// Use a pool with 1 worker and 1 queue slot, then fill both
+	p := &jobPools{
+		snmp:     newWorkerPool(4),
+		mikrotik: newWorkerPool(4),
+		ping:     newWorkerPool(4),
+		checks:   newWorkerPool(1),
+	}
+	t.Cleanup(func() { p.snmp.stop(); p.mikrotik.stop(); p.ping.stop(); p.checks.stop() })
+
+	done := make(chan struct{})
+	started := make(chan struct{})
+
+	// Block the single worker
+	p.checks.submit(context.Background(), func() {
+		close(started)
+		<-done
+	})
+	<-started
+
+	// Fill the queue slot
+	p.checks.submit(context.Background(), func() { <-done })
+
+	// Pool is now full — submit with cancelled ctx will fail
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	checkCh := make(chan *pb.CheckResult, 1)
+	check := &pb.Check{Id: "c1", CheckType: "tcp", TimeoutMs: 1000}
+	executeCheck(ctx, check, p, checkCh)
+	// Should log "check rejected (pool full)" but not panic
+	close(done)
+}
+
+func TestExecuteCheckCtxDoneInClosure(t *testing.T) {
+	// Tests the ctx.Done path inside executeCheck's closure:
+	// Submit succeeds, check runs, but result channel is blocked so ctx.Done fires.
+	p := testPools(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	checkCh := make(chan *pb.CheckResult) // unbuffered, no reader
+
+	check := &pb.Check{Id: "c1", CheckType: "tcp", TimeoutMs: 100,
+		Config: &pb.Check_Tcp{Tcp: &pb.TcpCheckConfig{Host: "127.0.0.1", Port: 1}}}
+
+	executeCheck(ctx, check, p, checkCh)
+
+	// Wait for the check to complete (TCP to port 1 fails fast)
+	time.Sleep(500 * time.Millisecond)
+	// Cancel ctx so the closure's select picks ctx.Done instead of blocked channel send
+	cancel()
+	time.Sleep(100 * time.Millisecond)
+}
+
+func TestRunSessionRestartInMainLoop(t *testing.T) {
+	srv := newFakeWSServer(t)
+
+	go func() {
+		srv.acceptAndJoin(t)
+		stopDrain := make(chan struct{})
+		go drainFrames(srv.conn, stopDrain)
+
+		time.Sleep(100 * time.Millisecond)
+		// Send restart event — exercised in the main loop select
+		srv.sendEvent("restart", json.RawMessage(`{}`))
+		time.Sleep(time.Second)
+		close(stopDrain)
+		srv.close()
+	}()
+
+	err := runSession(context.Background(), "ws://"+srv.addr(), "token")
+	if err != errRestartRequested {
+		t.Errorf("expected errRestartRequested, got: %v", err)
+	}
+}
+
+func TestRunSessionHeartbeats(t *testing.T) {
+	// Use short heartbeat intervals to exercise heartbeat paths
+	origHB := heartbeatInterval
+	origCHB := channelHeartbeatInterval
+	defer func() {
+		heartbeatInterval = origHB
+		channelHeartbeatInterval = origCHB
+	}()
+	heartbeatInterval = 100 * time.Millisecond
+	channelHeartbeatInterval = 100 * time.Millisecond
+
+	srv := newFakeWSServer(t)
+	go func() {
+		srv.acceptAndJoin(t)
+		stopDrain := make(chan struct{})
+		go drainFrames(srv.conn, stopDrain)
+		// Let heartbeats fire a few times
+		time.Sleep(500 * time.Millisecond)
+		close(stopDrain)
+		srv.close()
+	}()
+
+	err := runSession(context.Background(), "ws://"+srv.addr(), "token")
+	if err == nil {
+		t.Error("expected error after server close")
+	}
+}
+
+func TestRunAgentCancelDuringRetry(t *testing.T) {
+	// Connects to a port nothing listens on, then cancel during retry delay
+	ctx, cancel := context.WithCancel(context.Background())
+
+	done := make(chan struct{})
+	go func() {
+		runAgent(ctx, "ws://127.0.0.1:1", "token")
+		close(done)
+	}()
+
+	// Wait for first connection attempt to fail and retry delay to start
+	time.Sleep(1500 * time.Millisecond)
+	cancel()
+
+	select {
+	case <-done:
+		// runAgent returned after cancel during retry
+	case <-time.After(5 * time.Second):
+		t.Error("runAgent did not return after cancel during retry")
+	}
+}
diff --git a/checks_test.go b/checks_test.go
new file mode 100644
index 0000000..f2cd47d
--- /dev/null
+++ b/checks_test.go
@@ -0,0 +1,1372 @@
+package main
+
+import (
+	"bufio"
+	"context"
+	"crypto/tls"
+	"fmt"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/towerops-app/towerops-agent/pb"
+)
+
+// ---------------------------------------------------------------------------
+// ExecuteCheck routing tests
+// ---------------------------------------------------------------------------
+
+func TestExecuteCheck_UnknownCheckType(t *testing.T) {
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-1",
+		CheckType: "foobar",
+		TimeoutMs: 5000,
+	})
+	if result.Status != 3 {
+		t.Fatalf("expected status 3, got %d", result.Status)
+	}
+	if result.CheckId != "chk-1" {
+		t.Fatalf("expected CheckId chk-1, got %s", result.CheckId)
+	}
+	if !strings.Contains(result.Output, "Unknown check type") {
+		t.Fatalf("expected unknown check type message, got %s", result.Output)
+	}
+}
+
+func TestExecuteCheck_MissingHTTPConfig(t *testing.T) {
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-2",
+		CheckType: "http",
+		TimeoutMs: 5000,
+		// No Config set
+	})
+	if result.Status != 3 {
+		t.Fatalf("expected status 3, got %d", result.Status)
+	}
+	if !strings.Contains(result.Output, "Missing HTTP config") {
+		t.Fatalf("expected missing HTTP config message, got %s", result.Output)
+	}
+}
+
+func TestExecuteCheck_MissingTCPConfig(t *testing.T) {
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-3",
+		CheckType: "tcp",
+		TimeoutMs: 5000,
+	})
+	if result.Status != 3 {
+		t.Fatalf("expected status 3, got %d", result.Status)
+	}
+	if !strings.Contains(result.Output, "Missing TCP config") {
+		t.Fatalf("expected missing TCP config message, got %s", result.Output)
+	}
+}
+
+func TestExecuteCheck_MissingDNSConfig(t *testing.T) {
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-4",
+		CheckType: "dns",
+		TimeoutMs: 5000,
+	})
+	if result.Status != 3 {
+		t.Fatalf("expected status 3, got %d", result.Status)
+	}
+	if !strings.Contains(result.Output, "Missing DNS config") {
+		t.Fatalf("expected missing DNS config message, got %s", result.Output)
+	}
+}
+
+func TestExecuteCheck_SetsCheckIdAndTimestamp(t *testing.T) {
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-ts",
+		CheckType: "unknown",
+		TimeoutMs: 5000,
+	})
+	if result.CheckId != "chk-ts" {
+		t.Fatalf("expected CheckId chk-ts, got %s", result.CheckId)
+	}
+	if result.Timestamp == 0 {
+		t.Fatal("expected non-zero timestamp")
+	}
+	if result.ResponseTimeMs < 0 {
+		t.Fatalf("expected non-negative response time, got %f", result.ResponseTimeMs)
+	}
+}
+
+func TestExecuteCheck_HTTPRouting(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-route-http",
+		CheckType: "http",
+		TimeoutMs: 5000,
+		Config: &pb.Check_Http{
+			Http: &pb.HttpCheckConfig{
+				Url: srv.URL,
+			},
+		},
+	})
+	if result.Status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", result.Status, result.Output)
+	}
+}
+
+func TestExecuteCheck_TCPRouting(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	_, portStr, _ := net.SplitHostPort(ln.Addr().String())
+
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-route-tcp",
+		CheckType: "tcp",
+		TimeoutMs: 5000,
+		Config: &pb.Check_Tcp{
+			Tcp: &pb.TcpCheckConfig{
+				Host: "127.0.0.1",
+				Port: parsePort(portStr),
+			},
+		},
+	})
+	if result.Status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", result.Status, result.Output)
+	}
+}
+
+func TestExecuteCheck_DNSRouting(t *testing.T) {
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-route-dns",
+		CheckType: "dns",
+		TimeoutMs: 5000,
+		Config: &pb.Check_Dns{
+			Dns: &pb.DnsCheckConfig{
+				Hostname:   "localhost",
+				RecordType: "A",
+			},
+		},
+	})
+	// DNS for localhost may or may not resolve depending on system config,
+	// but the routing should work regardless.
+	if result.Status == 3 {
+		t.Fatalf("expected routing to DNS handler, got status 3 (UNKNOWN): %s", result.Output)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// HTTP check tests
+// ---------------------------------------------------------------------------
+
+func TestHTTPCheck_SuccessfulGET(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		fmt.Fprint(w, "OK")
+	}))
+	defer srv.Close()
+
+	status, output, rt := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: srv.URL,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "HTTP 200 OK") {
+		t.Fatalf("expected HTTP 200 OK in output, got %s", output)
+	}
+	if rt < 0 {
+		t.Fatalf("expected non-negative response time, got %f", rt)
+	}
+}
+
+func TestHTTPCheck_CustomMethod_POST(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != "POST" {
+			t.Errorf("expected POST, got %s", r.Method)
+			w.WriteHeader(405)
+			return
+		}
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:    srv.URL,
+		Method: "post",
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_CustomExpectedStatus(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(201)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:            srv.URL,
+		ExpectedStatus: 201,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_WrongStatusCode(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(500)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: srv.URL,
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "500") && !strings.Contains(output, "expected 200") {
+		t.Fatalf("expected status code mismatch message, got %s", output)
+	}
+}
+
+func TestHTTPCheck_DefaultMethodIsGET(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != "GET" {
+			t.Errorf("expected default GET, got %s", r.Method)
+			w.WriteHeader(405)
+			return
+		}
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:    srv.URL,
+		Method: "", // should default to GET
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_DefaultExpectedStatusIs200(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:            srv.URL,
+		ExpectedStatus: 0, // should default to 200
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_VerifySslFalseWithSelfSigned(t *testing.T) {
+	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	// With VerifySsl=false (InsecureSkipVerify=true), the self-signed cert should be accepted
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:       srv.URL,
+		VerifySsl: false,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 with VerifySsl=false, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_VerifySslTrueRejectsSelfSigned(t *testing.T) {
+	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	// With VerifySsl=true (InsecureSkipVerify=false), self-signed cert should fail
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:       srv.URL,
+		VerifySsl: true,
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2 with VerifySsl=true on self-signed, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Request failed") {
+		t.Fatalf("expected request failed message, got %s", output)
+	}
+}
+
+func TestHTTPCheck_FollowRedirectsTrue(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/redirect" {
+			http.Redirect(w, r, "/final", http.StatusMovedPermanently)
+			return
+		}
+		w.WriteHeader(200)
+		fmt.Fprint(w, "final page")
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:             srv.URL + "/redirect",
+		FollowRedirects: true,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 after following redirect, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "200") {
+		t.Fatalf("expected final 200 status in output, got %s", output)
+	}
+}
+
+func TestHTTPCheck_FollowRedirectsFalse(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == "/redirect" {
+			http.Redirect(w, r, "/final", http.StatusMovedPermanently)
+			return
+		}
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	// When FollowRedirects is false, we should see the 301 directly
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:             srv.URL + "/redirect",
+		FollowRedirects: false,
+		ExpectedStatus:  301,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 (matching 301), got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_FollowRedirectsFalse_DefaultExpects200(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, "/final", http.StatusMovedPermanently)
+	}))
+	defer srv.Close()
+
+	// FollowRedirects=false and default expected=200, but we get 301 → CRITICAL
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:             srv.URL + "/redirect",
+		FollowRedirects: false,
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2 (301 != 200), got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_RegexMatchSucceeds(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		fmt.Fprint(w, "Hello World! Version 1.2.3")
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:   srv.URL,
+		Regex: `Version \d+\.\d+\.\d+`,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 for regex match, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_RegexMatchFails(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		fmt.Fprint(w, "Hello World!")
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:   srv.URL,
+		Regex: `Version \d+`,
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2 for regex mismatch, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "does not match") {
+		t.Fatalf("expected 'does not match' in output, got %s", output)
+	}
+}
+
+func TestHTTPCheck_InvalidRegex(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		fmt.Fprint(w, "some body")
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:   srv.URL,
+		Regex: `[invalid`,
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2 for invalid regex, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Invalid regex") {
+		t.Fatalf("expected 'Invalid regex' in output, got %s", output)
+	}
+}
+
+func TestHTTPCheck_CustomHeaders(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Header.Get("X-Custom") != "test-value" {
+			t.Errorf("expected X-Custom header to be test-value, got %s", r.Header.Get("X-Custom"))
+			w.WriteHeader(400)
+			return
+		}
+		if r.Header.Get("Authorization") != "Bearer abc123" {
+			t.Errorf("expected Authorization header, got %s", r.Header.Get("Authorization"))
+			w.WriteHeader(401)
+			return
+		}
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: srv.URL,
+		Headers: map[string]string{
+			"X-Custom":      "test-value",
+			"Authorization": "Bearer abc123",
+		},
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_UnreachableServer(t *testing.T) {
+	// Use a non-routable address to guarantee failure
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: "http://192.0.2.1:1/unreachable",
+	}, 1000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Request failed") {
+		t.Fatalf("expected 'Request failed' in output, got %s", output)
+	}
+}
+
+func TestHTTPCheck_InvalidURL(t *testing.T) {
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: "://not-a-url",
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Failed to create request") {
+		t.Fatalf("expected 'Failed to create request' in output, got %s", output)
+	}
+}
+
+func TestHTTPCheck_RequestWithBody(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		buf := make([]byte, 1024)
+		n, _ := r.Body.Read(buf)
+		body := string(buf[:n])
+		if body != `{"key":"value"}` {
+			t.Errorf("expected JSON body, got %s", body)
+			w.WriteHeader(400)
+			return
+		}
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:    srv.URL,
+		Method: "POST",
+		Body:   `{"key":"value"}`,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_ContextCancellation(t *testing.T) {
+	// Server that hangs forever
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		<-r.Context().Done()
+	}))
+	defer srv.Close()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	// Cancel after a short delay
+	go func() {
+		time.Sleep(50 * time.Millisecond)
+		cancel()
+	}()
+
+	status, output, _ := executeHTTPCheck(ctx, &pb.HttpCheckConfig{
+		Url: srv.URL,
+	}, 30000) // long timeout so the context cancel hits first
+
+	if status != 2 {
+		t.Fatalf("expected status 2 on context cancel, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Request failed") {
+		t.Fatalf("expected 'Request failed' in output, got %s", output)
+	}
+}
+
+func TestHTTPCheck_LargeResponseBodyWithRegex(t *testing.T) {
+	// Generate a large body (500KB) with a marker at the end
+	bigBody := strings.Repeat("a", 500*1024) + "MARKER_FOUND_HERE"
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		fmt.Fprint(w, bigBody)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:   srv.URL,
+		Regex: `MARKER_FOUND_HERE`,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 for regex match in large body, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_SlowServerTimeout(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(500 * time.Millisecond)
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: srv.URL,
+	}, 50) // 50ms timeout, server takes 500ms
+
+	if status != 2 {
+		t.Fatalf("expected status 2 for timeout, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_EmptyBody_NoRegex(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: srv.URL,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// TCP check tests
+// ---------------------------------------------------------------------------
+
+func TestTCPCheck_PortOpen(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	// Accept connections in background so dial doesn't hang
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			conn.Close()
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	status, output, rt := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host: "127.0.0.1",
+		Port: port,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, fmt.Sprintf("TCP port %d open", port)) {
+		t.Fatalf("expected port open message, got %s", output)
+	}
+	if rt < 0 {
+		t.Fatalf("expected non-negative response time, got %f", rt)
+	}
+}
+
+func TestTCPCheck_PortClosed(t *testing.T) {
+	// Find a port that's definitely not listening by binding and immediately closing
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	port := parsePort(portFromListener(ln))
+	ln.Close() // close immediately so port is refused
+
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host: "127.0.0.1",
+		Port: port,
+	}, 2000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Connection failed") {
+		t.Fatalf("expected 'Connection failed' in output, got %s", output)
+	}
+}
+
+func TestTCPCheck_SendExpectSuccess(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	// Echo server
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				scanner := bufio.NewScanner(c)
+				if scanner.Scan() {
+					line := scanner.Text()
+					fmt.Fprintf(c, "ECHO:%s\n", line)
+				}
+			}(conn)
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host:   "127.0.0.1",
+		Port:   port,
+		Send:   "hello\n",
+		Expect: "ECHO:hello",
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestTCPCheck_SendExpectMismatch(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				buf := make([]byte, 1024)
+				c.Read(buf)
+				fmt.Fprint(c, "WRONG_RESPONSE")
+			}(conn)
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host:   "127.0.0.1",
+		Port:   port,
+		Send:   "hello",
+		Expect: "EXPECTED_VALUE",
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Unexpected response") {
+		t.Fatalf("expected 'Unexpected response' in output, got %s", output)
+	}
+}
+
+func TestTCPCheck_SendWithEmptyExpect(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				buf := make([]byte, 1024)
+				c.Read(buf)
+				// Don't send anything back
+			}(conn)
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host:   "127.0.0.1",
+		Port:   port,
+		Send:   "data\n",
+		Expect: "", // no expect check
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 (send only, no expect), got %d: %s", status, output)
+	}
+}
+
+func TestTCPCheck_IPv6Localhost(t *testing.T) {
+	ln, err := net.Listen("tcp6", "[::1]:0")
+	if err != nil {
+		t.Skip("IPv6 not available on this system")
+	}
+	defer ln.Close()
+
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			conn.Close()
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host: "::1",
+		Port: port,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 for IPv6, got %d: %s", status, output)
+	}
+}
+
+func TestTCPCheck_ReadTimeoutOnExpect(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	// Server that accepts and reads but never writes back
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				buf := make([]byte, 1024)
+				c.Read(buf)
+				// Intentionally never respond - hold connection open
+				time.Sleep(10 * time.Second)
+			}(conn)
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host:   "127.0.0.1",
+		Port:   port,
+		Send:   "hello",
+		Expect: "response",
+	}, 200) // short timeout
+
+	if status != 2 {
+		t.Fatalf("expected status 2 for read timeout, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Receive failed") {
+		t.Fatalf("expected 'Receive failed' in output, got %s", output)
+	}
+}
+
+func TestTCPCheck_VeryShortTimeout(t *testing.T) {
+	// Use TEST-NET address that won't respond
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host: "192.0.2.1",
+		Port: 80,
+	}, 1) // 1ms timeout - should fail
+
+	if status != 2 {
+		t.Fatalf("expected status 2, got %d: %s", status, output)
+	}
+}
+
+func TestTCPCheck_BinaryDataSendExpect(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	// Server that echoes binary data back
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer c.Close()
+				buf := make([]byte, 4096)
+				n, err := c.Read(buf)
+				if err != nil {
+					return
+				}
+				c.Write(buf[:n])
+			}(conn)
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	// Send some binary-ish data
+	sendData := "BIN\x00\x01\x02DATA"
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host:   "127.0.0.1",
+		Port:   port,
+		Send:   sendData,
+		Expect: sendData,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 for binary echo, got %d: %s", status, output)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// DNS check tests
+// ---------------------------------------------------------------------------
+
+func TestDNSCheck_ARecord(t *testing.T) {
+	status, output, rt := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "localhost",
+		RecordType: "A",
+	}, 5000)
+
+	// localhost should resolve on most systems, but skip if it doesn't
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available for localhost")
+	}
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Resolved to:") {
+		t.Fatalf("expected 'Resolved to:' in output, got %s", output)
+	}
+	if rt < 0 {
+		t.Fatalf("expected non-negative response time, got %f", rt)
+	}
+}
+
+func TestDNSCheck_AAAARecord(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "localhost",
+		RecordType: "AAAA",
+	}, 5000)
+
+	// AAAA for localhost may or may not exist - just verify it doesn't crash
+	// and returns a valid status
+	if status != 0 && status != 2 {
+		t.Fatalf("expected status 0 or 2, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_CNAMERecord(t *testing.T) {
+	// Use a well-known domain that has a CNAME
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "www.google.com",
+		RecordType: "CNAME",
+	}, 5000)
+
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available")
+	}
+	// CNAME lookup may return the hostname itself if no CNAME exists
+	if status != 0 && status != 2 {
+		t.Fatalf("expected status 0 or 2, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_MXRecord(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "google.com",
+		RecordType: "MX",
+	}, 5000)
+
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available")
+	}
+	if status != 0 {
+		t.Fatalf("expected status 0 for MX lookup, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Resolved to:") {
+		t.Fatalf("expected 'Resolved to:' in output, got %s", output)
+	}
+}
+
+func TestDNSCheck_TXTRecord(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "google.com",
+		RecordType: "TXT",
+	}, 5000)
+
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available")
+	}
+	if status != 0 {
+		t.Fatalf("expected status 0 for TXT lookup, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_UnsupportedRecordType(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "example.com",
+		RecordType: "SRV",
+	}, 5000)
+
+	if status != 3 {
+		t.Fatalf("expected status 3, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Unsupported record type") {
+		t.Fatalf("expected 'Unsupported record type' in output, got %s", output)
+	}
+}
+
+func TestDNSCheck_ExpectedMatchesResult(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "localhost",
+		RecordType: "A",
+		Expected:   "127.0.0.1",
+	}, 5000)
+
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available for localhost")
+	}
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_ExpectedDoesNotMatch(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "localhost",
+		RecordType: "A",
+		Expected:   "10.99.99.99",
+	}, 5000)
+
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available for localhost")
+	}
+	// If localhost resolved, the expected won't match
+	if status != 2 {
+		t.Fatalf("expected status 2 for mismatched expected, got %d: %s", status, output)
+	}
+	if !strings.Contains(output, "Expected '10.99.99.99'") {
+		t.Fatalf("expected mismatch message, got %s", output)
+	}
+}
+
+func TestDNSCheck_NonexistentDomain(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "this-domain-does-not-exist-towerops-test.invalid",
+		RecordType: "A",
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2 for NXDOMAIN, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_DefaultRecordTypeIsA(t *testing.T) {
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "localhost",
+		RecordType: "", // should default to A
+	}, 5000)
+
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available for localhost")
+	}
+	// Should behave the same as explicitly specifying "A"
+	if status != 0 {
+		t.Fatalf("expected status 0 with default record type, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_CustomDNSServer(t *testing.T) {
+	// Test with Google's public DNS
+	conn, err := net.DialTimeout("udp", "8.8.8.8:53", 2*time.Second)
+	if err != nil {
+		t.Skip("Cannot reach 8.8.8.8:53, skipping custom DNS server test")
+	}
+	conn.Close()
+
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "example.com",
+		RecordType: "A",
+		Server:     "8.8.8.8",
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 with custom DNS server, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_RecordTypeCaseInsensitive(t *testing.T) {
+	// The code does strings.ToUpper, so lowercase should work
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "localhost",
+		RecordType: "a", // lowercase
+	}, 5000)
+
+	if status == 2 && strings.Contains(output, "DNS query failed") {
+		t.Skip("DNS resolution not available for localhost")
+	}
+	if status != 0 {
+		t.Fatalf("expected status 0 with lowercase record type, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_NoRecordsFound(t *testing.T) {
+	// AAAA for a domain that likely only has A records
+	// Use a known domain that almost certainly won't have AAAA
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "this-domain-does-not-exist-towerops-test.invalid",
+		RecordType: "AAAA",
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2, got %d: %s", status, output)
+	}
+}
+
+func TestDNSCheck_VeryShortTimeout(t *testing.T) {
+	// Use a custom DNS server with 1ms timeout - should timeout
+	status, output, _ := executeDNSCheck(context.Background(), &pb.DnsCheckConfig{
+		Hostname:   "example.com",
+		RecordType: "A",
+		Server:     "8.8.8.8",
+	}, 1) // 1ms timeout
+
+	// Should fail due to timeout
+	if status != 2 {
+		// On very fast networks this could actually succeed, so just verify it doesn't crash
+		t.Logf("DNS with 1ms timeout returned status %d: %s (may succeed on fast networks)", status, output)
+	}
+}
+
+func TestTCPCheck_SendFailsOnClosedConnection(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	// Server accepts and immediately closes the connection
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			conn.Close()
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+
+	// Send a large payload to trigger a write error on a closed connection.
+	// The first small write might succeed (kernel buffer), but a large write
+	// after the peer has closed should fail with a broken pipe or similar.
+	largePayload := strings.Repeat("x", 1024*1024) // 1MB
+	// Give the server time to close the connection
+	time.Sleep(50 * time.Millisecond)
+
+	status, output, _ := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host:   "127.0.0.1",
+		Port:   port,
+		Send:   largePayload,
+		Expect: "something",
+	}, 5000)
+
+	// This may hit either "Send failed" or "Receive failed" depending on timing
+	if status != 2 {
+		t.Fatalf("expected status 2 for write to closed conn, got %d: %s", status, output)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// Edge case tests
+// ---------------------------------------------------------------------------
+
+func TestHTTPCheck_MethodUppercased(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != "PUT" {
+			t.Errorf("expected PUT, got %s", r.Method)
+			w.WriteHeader(405)
+			return
+		}
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:    srv.URL,
+		Method: "put", // lowercase
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_ResponseTimeIsPositive(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		time.Sleep(10 * time.Millisecond)
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	_, _, rt := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url: srv.URL,
+	}, 5000)
+
+	if rt <= 0 {
+		t.Fatalf("expected positive response time, got %f", rt)
+	}
+}
+
+func TestTCPCheck_ResponseTimeReported(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ln.Close()
+
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			conn.Close()
+		}
+	}()
+
+	port := parsePort(portFromListener(ln))
+	_, _, rt := executeTCPCheck(context.Background(), &pb.TcpCheckConfig{
+		Host: "127.0.0.1",
+		Port: port,
+	}, 5000)
+
+	if rt < 0 {
+		t.Fatalf("expected non-negative response time, got %f", rt)
+	}
+}
+
+func TestExecuteCheck_ResponseTimeFallback(t *testing.T) {
+	// When status/output are set directly (unknown type), responseTimeMs should be calculated
+	result := ExecuteCheck(context.Background(), &pb.Check{
+		Id:        "chk-fallback",
+		CheckType: "invalid",
+		TimeoutMs: 5000,
+	})
+
+	if result.ResponseTimeMs < 0 {
+		t.Fatalf("expected non-negative response time, got %f", result.ResponseTimeMs)
+	}
+}
+
+func TestHTTPCheck_TLSServerNoVerify(t *testing.T) {
+	// Create a TLS server with custom cert
+	srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		fmt.Fprint(w, "secure")
+	}))
+	srv.TLS = &tls.Config{}
+	srv.StartTLS()
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:       srv.URL,
+		VerifySsl: false,
+		Regex:     "secure",
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_HeadMethod(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != "HEAD" {
+			t.Errorf("expected HEAD, got %s", r.Method)
+		}
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:    srv.URL,
+		Method: "HEAD",
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_DeleteMethod(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != "DELETE" {
+			t.Errorf("expected DELETE, got %s", r.Method)
+		}
+		w.WriteHeader(204)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:            srv.URL,
+		Method:         "DELETE",
+		ExpectedStatus: 204,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_RegexOnEmptyBody(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+		// No body written
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:   srv.URL,
+		Regex: "something",
+	}, 5000)
+
+	if status != 2 {
+		t.Fatalf("expected status 2 (regex no match on empty body), got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_MultipleRedirects(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case "/a":
+			http.Redirect(w, r, "/b", http.StatusFound)
+		case "/b":
+			http.Redirect(w, r, "/c", http.StatusFound)
+		case "/c":
+			w.WriteHeader(200)
+			fmt.Fprint(w, "final")
+		}
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:             srv.URL + "/a",
+		FollowRedirects: true,
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 after multiple redirects, got %d: %s", status, output)
+	}
+}
+
+func TestHTTPCheck_EmptyHeadersMap(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	}))
+	defer srv.Close()
+
+	status, output, _ := executeHTTPCheck(context.Background(), &pb.HttpCheckConfig{
+		Url:     srv.URL,
+		Headers: map[string]string{},
+	}, 5000)
+
+	if status != 0 {
+		t.Fatalf("expected status 0 with empty headers, got %d: %s", status, output)
+	}
+}
+
+// ---------------------------------------------------------------------------
+// Helper functions
+// ---------------------------------------------------------------------------
+
+// portFromListener extracts the port string from a net.Listener's address.
+func portFromListener(ln net.Listener) string {
+	_, port, _ := net.SplitHostPort(ln.Addr().String())
+	return port
+}
+
+// parsePort converts a port string to uint32 for use in proto configs.
+func parsePort(s string) uint32 {
+	var port uint32
+	fmt.Sscanf(s, "%d", &port)
+	return port
+}
diff --git a/hostkeys_test.go b/hostkeys_test.go
index b36586c..7ecc781 100644
--- a/hostkeys_test.go
+++ b/hostkeys_test.go
@@ -1,7 +1,7 @@
 package main
 
 import (
-	"os"
+	"crypto/x509"
 	"path/filepath"
 	"sync"
 	"testing"
@@ -116,7 +116,20 @@ func TestHostKeyStoreSaveError(t *testing.T) {
 }
 
 func TestTlsCertFingerprint(t *testing.T) {
-	// Just verify it doesn't panic with nil - we test with real certs in integration
-	// The function is simple enough: sha256 of cert.Raw
-	_ = os.Getenv("dummy") // placeholder
+	cert := &x509.Certificate{
+		Raw: []byte("test certificate data"),
+	}
+	fp := tlsCertFingerprint(cert)
+	if fp == "" {
+		t.Error("expected non-empty fingerprint")
+	}
+	// Verify it's a hex-encoded SHA256 (64 chars)
+	if len(fp) != 64 {
+		t.Errorf("expected 64-char hex fingerprint, got %d chars: %s", len(fp), fp)
+	}
+	// Same input should produce same output
+	fp2 := tlsCertFingerprint(cert)
+	if fp != fp2 {
+		t.Error("fingerprint not deterministic")
+	}
 }
diff --git a/main.go b/main.go
index af71359..bd9a258 100644
--- a/main.go
+++ b/main.go
@@ -93,7 +93,7 @@ func toWebSocketURL(rawURL string) string {
 
 	if strings.HasPrefix(result, "ws://") && !insecureFlag {
 		slog.Error("plaintext ws:// connection rejected — use wss:// or pass --insecure to allow")
-		os.Exit(1)
+		osExit(1)
 	}
 	return result
 }
diff --git a/main_test.go b/main_test.go
index bb7b7ab..3b57722 100644
--- a/main_test.go
+++ b/main_test.go
@@ -39,6 +39,14 @@ func TestSanitizeURL(t *testing.T) {
 	}
 }
 
+func TestIsFlagSet(t *testing.T) {
+	// In tests, flags are not set via flag.Parse on the default flag set
+	// so isFlagSet should return false for any arbitrary name.
+	if isFlagSet("nonexistent-flag-xyz") {
+		t.Error("expected isFlagSet to return false for unset flag")
+	}
+}
+
 func TestToWebSocketURL(t *testing.T) {
 	// Enable insecure for testing plaintext conversions
 	origInsecure := insecureFlag
@@ -62,3 +70,33 @@ func TestToWebSocketURL(t *testing.T) {
 		}
 	}
 }
+
+func TestToWebSocketURLRejectsPlaintext(t *testing.T) {
+	origInsecure := insecureFlag
+	origExit := osExit
+	defer func() {
+		insecureFlag = origInsecure
+		osExit = origExit
+	}()
+
+	insecureFlag = false
+	exitCode := -1
+	osExit = func(code int) { exitCode = code }
+
+	tests := []struct {
+		name  string
+		input string
+	}{
+		{"ws:// scheme", "ws://localhost:4000"},
+		{"http:// converts to ws://", "http://localhost:4000"},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			exitCode = -1
+			toWebSocketURL(tt.input)
+			if exitCode != 1 {
+				t.Errorf("expected osExit(1), got %d", exitCode)
+			}
+		})
+	}
+}
diff --git a/mikrotik_test.go b/mikrotik_test.go
index bc8105e..1872287 100644
--- a/mikrotik_test.go
+++ b/mikrotik_test.go
@@ -2,10 +2,21 @@ package main
 
 import (
 	"bytes"
+	"crypto/ecdsa"
+	"crypto/elliptic"
+	"crypto/rand"
+	"crypto/tls"
+	"crypto/x509"
+	"encoding/pem"
 	"fmt"
 	"io"
+	"math/big"
 	"net"
+	"path/filepath"
+	"strings"
+	"sync"
 	"testing"
+	"time"
 )
 
 func TestEncodeLength(t *testing.T) {
@@ -467,6 +478,78 @@ func TestMikrotikConnectSSL(t *testing.T) {
 	}
 }
 
+func TestMikrotikConnectSSLWithServer(t *testing.T) {
+	// Generate a self-signed cert at runtime
+	cert := generateTestCert(t)
+	tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}}
+	ln, err := tls.Listen("tcp", "127.0.0.1:0", tlsConfig)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	// Reset host key store for this test
+	origStore := globalHostKeys
+	defer func() {
+		hostKeysOnce = sync.Once{}
+		globalHostKeys = origStore
+	}()
+	hostKeysOnce = sync.Once{}
+	t.Setenv("TOWEROPS_HOST_KEYS_FILE", filepath.Join(t.TempDir(), "hosts.json"))
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+		sc := &mikrotikClient{conn: conn}
+		_, _ = sc.readSentence()
+		_ = sc.writeSentence([]string{"!done"})
+		_, _ = sc.readSentence()
+		_ = sc.writeSentence([]string{"!fatal"})
+	}()
+
+	_, port, _ := net.SplitHostPort(ln.Addr().String())
+	var portNum uint32
+	_, _ = fmt.Sscanf(port, "%d", &portNum)
+
+	client, err := mikrotikConnect("127.0.0.1", portNum, "admin", "pass", true)
+	if err != nil {
+		t.Fatalf("expected TLS connection to succeed: %v", err)
+	}
+	_ = client.close()
+}
+
+func generateTestCert(t *testing.T) tls.Certificate {
+	t.Helper()
+	key, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
+	if err != nil {
+		t.Fatal(err)
+	}
+	tmpl := &x509.Certificate{
+		SerialNumber: big.NewInt(1),
+		NotBefore:    time.Now(),
+		NotAfter:     time.Now().Add(time.Hour),
+		IPAddresses:  []net.IP{net.ParseIP("127.0.0.1")},
+	}
+	certDER, err := x509.CreateCertificate(rand.Reader, tmpl, tmpl, &key.PublicKey, key)
+	if err != nil {
+		t.Fatal(err)
+	}
+	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
+	keyDER, err := x509.MarshalECPrivateKey(key)
+	if err != nil {
+		t.Fatal(err)
+	}
+	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyDER})
+	tlsCert, err := tls.X509KeyPair(certPEM, keyPEM)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return tlsCert
+}
+
 func TestReadResponseEmptySentence(t *testing.T) {
 	// An empty sentence (just the terminator byte) should be skipped
 	var buf bytes.Buffer
@@ -636,3 +719,51 @@ func (rwc *readWriteCloser) Close() error {
 	}
 	return nil
 }
+
+func TestMikrotikConnectSSLTOFUMismatch(t *testing.T) {
+	cert := generateTestCert(t)
+	tlsConfig := &tls.Config{Certificates: []tls.Certificate{cert}}
+	ln, err := tls.Listen("tcp", "127.0.0.1:0", tlsConfig)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	// Reset host key store and pre-populate with wrong fingerprint
+	origStore := globalHostKeys
+	defer func() {
+		hostKeysOnce = sync.Once{}
+		globalHostKeys = origStore
+	}()
+	hostKeysOnce = sync.Once{}
+	t.Setenv("TOWEROPS_HOST_KEYS_FILE", filepath.Join(t.TempDir(), "hosts.json"))
+
+	_, port, _ := net.SplitHostPort(ln.Addr().String())
+	var portNum uint32
+	_, _ = fmt.Sscanf(port, "%d", &portNum)
+
+	// Pre-register wrong fingerprint so TOFU verification fails
+	store := getHostKeyStore()
+	store.keys["tls:"+net.JoinHostPort("127.0.0.1", port)] = "wrong_fingerprint"
+
+	go func() {
+		conn, err := ln.Accept()
+		if err != nil {
+			return
+		}
+		defer func() { _ = conn.Close() }()
+		// Complete TLS handshake so client gets peer certificates
+		if tlsConn, ok := conn.(*tls.Conn); ok {
+			_ = tlsConn.Handshake()
+		}
+		time.Sleep(time.Second)
+	}()
+
+	_, err = mikrotikConnect("127.0.0.1", portNum, "admin", "pass", true)
+	if err == nil {
+		t.Error("expected TOFU verification failure")
+	}
+	if !strings.Contains(err.Error(), "TOFU verification failed") {
+		t.Errorf("expected 'TOFU verification failed' in error, got: %v", err)
+	}
+}
diff --git a/ping_test.go b/ping_test.go
index e5b2fb5..5a91d4d 100644
--- a/ping_test.go
+++ b/ping_test.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"fmt"
 	"runtime"
 	"testing"
 )
@@ -65,6 +66,13 @@ func TestIcmpPingInvalidIP(t *testing.T) {
 	}
 }
 
+func TestErrICMPUnavailableError(t *testing.T) {
+	err := &errICMPUnavailable{err: fmt.Errorf("permission denied")}
+	if err.Error() != "permission denied" {
+		t.Errorf("got %q, want %q", err.Error(), "permission denied")
+	}
+}
+
 func TestParsePingTime(t *testing.T) {
 	tests := []struct {
 		name    string
diff --git a/websocket_test.go b/websocket_test.go
index 659cea3..7bf9410 100644
--- a/websocket_test.go
+++ b/websocket_test.go
@@ -657,6 +657,60 @@ func TestWSDialIPv4Fallback(t *testing.T) {
 	}
 }
 
+func TestWSDialTLSPath(t *testing.T) {
+	// Verify the TLS dial path is exercised (connection will fail, but we hit the code path)
+	origTLS := tlsDial
+	defer func() { tlsDial = origTLS }()
+
+	called := false
+	tlsDial = func(network, addr string) (net.Conn, error) {
+		called = true
+		return nil, fmt.Errorf("tls dial: connection refused")
+	}
+
+	_, err := WSDial("wss://127.0.0.1:9999/path")
+	if err == nil {
+		t.Error("expected TLS dial error")
+	}
+	if !called {
+		t.Error("expected tlsDial to be called for wss:// URL")
+	}
+}
+
+func TestWSDialMissingAcceptHeader(t *testing.T) {
+	// Server sends 101 but without Sec-WebSocket-Accept header
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer func() { _ = ln.Close() }()
+
+	go func() {
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return
+			}
+			go func(c net.Conn) {
+				defer func() { _ = c.Close() }()
+				buf := make([]byte, 4096)
+				_, _ = c.Read(buf)
+				resp := "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\nConnection: Upgrade\r\n\r\n"
+				_, _ = c.Write([]byte(resp))
+			}(conn)
+		}
+	}()
+
+	addr := ln.Addr().String()
+	_, err = WSDial("ws://" + addr + "/socket")
+	if err == nil {
+		t.Error("expected error for missing accept header")
+	}
+	if !strings.Contains(err.Error(), "missing Sec-WebSocket-Accept") {
+		t.Errorf("expected 'missing Sec-WebSocket-Accept' in error, got: %v", err)
+	}
+}
+
 func TestWSDialDefaultPorts(t *testing.T) {
 	// Test that ws:// defaults to port 80 — will fail to connect but verifies URL parsing
 	_, err := WSDial("ws://127.0.0.1/path")