Locust 是一個強大的負載測試工具，但它的核心客戶端主要是為 HTTP/HTTPS 協定所設計。然而，透過 Custom Clients (自訂客戶端) 功能，你可以擴展 Locust 的能力，讓它能夠對任何協定進行負載測試，例如 FTP、TCP、WebSocket，甚至是自訂的二進位協定。

這個功能的核心概念是：你來處理底層的網路連線和資料傳輸，然後將每次操作的結果（例如請求類型、回應時間、回應大小等）回報給 Locust，讓 Locust 的統計系統能夠統一收集和呈現這些資料。

基本概念

自訂客戶端模式讓我們能夠：

• 測試非 HTTP 協定：將 Locust 應用於檔案伺服器、即時通訊服務或遊戲伺服器等。
• 重用 Locust 功能：享受 Locust 強大的負載生成、即時統計和報告功能，無需從頭開發。
• 保持分散式測試能力：自訂客戶端完全支援 Locust 的分散式模式，讓你能夠在多台機器上擴大測試規模。

專案目錄結構

一個典型的 Locust 自訂客戶端專案會將客戶端實作與主測試腳本分開，以保持程式碼的清晰和可維護性。

project/
├── locustfile.py           # 主要的測試腳本
├── clients/                # 專門存放自訂客戶端的資料夾
│   ├── __init__.py         # 讓 Python 識別 clients 為套件
│   ├── ftp_client.py       # FTP 客戶端的實作
└── requirements.txt        # 專案所需的套件

FTP Client 範例

在這個範例中，我們將為 FTP（File Transfer Protocol）建立一個自訂客戶端，並模擬多個使用者同時上傳檔案的負載情境。

1. 實作 FTP 客戶端 (clients/ftp_client.py)

首先，我們需要一個類別來封裝所有的 FTP 相關操作。我們將使用 Python 內建的 ftplib 模組來實現。

FTPClient 類別：
這個類別負責處理底層的 FTP 連線和操作。每個方法都包含時間測量和事件回報的邏輯。

connect() 方法：
負責連線和登入 FTP 伺服器。events.request.fire() 方法是關鍵，它會將這次「請求」的資料回傳給 Locust。

• request_type: 自訂的請求類型，例如 "FTP"。
• name: 請求的名稱，例如 "connect" 或 "upload"。
• response_time: 請求所需時間，以毫秒為單位。
• response_length: 回應的資料大小，單位為位元組。
• exception: 如果發生錯誤，傳遞錯誤物件，Locust 會將其標記為失敗。

FTPUser 類別：
這個類別繼承自 Locust 的 User 類別，並將 abstract 設為 True，表示它不會被直接執行，而是作為其他測試類別的基底。

• on_start() 和 on_stop()：Locust 會在每個模擬使用者（User）開始和停止時分別呼叫這兩個方法。我們在這裡處理 FTP 連線的建立和斷開，確保每個使用者在測試期間都維持一個獨立的連線。

# clients/ftp_client.py
import ftplib
import time
import os
from locust import User, events
from locust.exception import LocustError

class FTPClient:
    """
    FTP 客戶端，用於處理底層的 FTP 連線和操作。
    """
    def __init__(self, host, port=21):
        self.host = host
        self.port = port
        self.ftp = None
    
    def connect(self, username='anonymous', password=''):
        start_time = time.time()
        try:
            self.ftp = ftplib.FTP()
            self.ftp.connect(self.host, self.port)
            self.ftp.login(username, password)
            
            events.request.fire(
                request_type="FTP",
                name="connect",
                response_time=int((time.time() - start_time) * 1000),
                response_length=0,
            )
            return True
        except Exception as e:
            events.request.fire(
                request_type="FTP",
                name="connect",
                response_time=int((time.time() - start_time) * 1000),
                response_length=0,
                exception=e
            )
            raise LocustError(f"FTP 連接失敗: {e}")
    
    def upload_file(self, local_file_path, remote_filename):
        start_time = time.time()
        try:
            with open(local_file_path, 'rb') as file:
                self.ftp.storbinary(f'STOR {remote_filename}', file)
            
            file_size = os.path.getsize(local_file_path)
            events.request.fire(
                request_type="FTP",
                name="upload",
                response_time=int((time.time() - start_time) * 1000),
                response_length=file_size,
            )
            return True
        except Exception as e:
            events.request.fire(
                request_type="FTP",
                name="upload",
                response_time=int((time.time() - start_time) * 1000),
                response_length=0,
                exception=e
            )
            raise LocustError(f"FTP 上傳失敗: {e}")
    
    def disconnect(self):
        if self.ftp:
            try:
                self.ftp.quit()
            except:
                self.ftp.close()

class FTPUser(User):
    abstract = True
    
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.client = FTPClient(self.host)
    
    def on_start(self):
        self.client.connect()
    
    def on_stop(self):
        self.client.disconnect()

2. 使用 FTP 客戶端 (locustfile.py)

在主測試腳本中，我們只需要簡單地繼承 FTPUser 類別，並定義具體的測試任務。

# locustfile.py
import os
import random
import tempfile
from locust import task, between
from clients.ftp_client import FTPUser

class FTPLoadTest(FTPUser):
    # 每個使用者任務之間的等待時間為 1 到 3 秒
    wait_time = between(1, 3)
    
    # 指定要測試的 FTP 伺服器
    # ⚠️ 請替換成你的測試伺服器位址！
    host = "ftp.example.com" 
    
    @task
    def test_upload(self):
        # 1. 在本地建立一個臨時測試檔案
        local_file = tempfile.NamedTemporaryFile(delete=False)
        local_file.write(f"Test content: {random.randint(1, 1000)}".encode('utf-8'))
        local_file.close()
        
        try:
            # 2. 呼叫自訂客戶端的方法來上傳檔案
            remote_name = f"uploaded_file_{os.path.basename(local_file.name)}.txt"
            self.client.upload_file(local_file.name, remote_name)
        finally:
            # 3. 測試完成後清理臨時檔案
            os.remove(local_file.name)

其他 Client 範例

Locust 的自訂客戶端模式具有高度的彈性，可以套用到任何協定上。以下是其他幾個常見協定的簡要範例。

TCP Client

TCP 是許多低層次服務的基礎。你可以使用 Python 內建的 socket 模組來實作一個 TCP 客戶端。

• TCPClient: 負責建立 socket 連線並發送資料。
• TCPUser: 繼承 User，管理每個使用者與伺服器的 TCP 持久連線。

WebSocket Client

WebSocket 常用於即時通訊，例如聊天室或即時數據推送。我們可以使用 websocket-client 等第三方套件來實作。

• WebSocketClient: 處理 WebSocket 連線、發送訊息和斷開連線。
• WebSocketUser: 繼承 User，確保每個使用者在測試期間都維持一個 WebSocket 連線。

執行測試

一旦你完成了客戶端和測試腳本的撰寫，就可以用 locust 命令來啟動測試。

1. 安裝依賴：如果你的客戶端需要額外的套件，請先安裝。例如，對於 WebSocket，你需要 pip install locust websocket-client。

2. 執行命令：

   • 針對 FTP 測試：locust -f locustfile.py --host ftp://ftp.example.com
   • 針對 TCP 測試：locust -f tcp_test.py --host localhost:8888
   • 針對 WebSocket 測試：locust -f ws_test.py --host ws://localhost:8080/ws

Locust 會自動啟動 Web 介面，你可以在上面看到所有請求的效能統計，包括你自訂的 FTP、TCP 或 WebSocket 請求。

總結

透過 Custom Clients，Locust 的應用範圍從單純的 Web 測試擴展到了任何需要效能驗證的網路服務。這種模式讓你能夠利用 Locust 成熟的框架來專注於測試邏輯本身，而無需擔心底層的負載生成和結果統計，大大提升了測試的效率和靈活性。