mirror of
https://github.com/logos-storage/bittorrent-benchmarks.git
synced 2026-01-03 05:23:06 +00:00
75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
from typing import Dict, Optional, IO
|
|
|
|
from aiohttp import ClientTimeout
|
|
|
|
from benchmarks.codex.client.async_client import AsyncCodexClient, Cid, DownloadStatus
|
|
from benchmarks.codex.client.common import Manifest
|
|
|
|
|
|
class FakeDownload:
|
|
def __init__(self, manifest: Manifest) -> None:
|
|
self.manifest = manifest
|
|
self.downloaded = 0
|
|
self.exception: Optional[Exception] = None
|
|
|
|
def advance_download(self, blocks: int):
|
|
self.downloaded += blocks
|
|
print("Advance download to", self.downloaded)
|
|
assert (
|
|
self.downloaded <= self.manifest.block_count
|
|
), "Downloaded blocks exceed total blocks"
|
|
|
|
def abort(self, exception: Exception):
|
|
self.exception = exception
|
|
|
|
|
|
class FakeCodex(AsyncCodexClient):
|
|
def __init__(self) -> None:
|
|
self.manifests: Dict[Cid, Manifest] = {}
|
|
self.ongoing_downloads: Dict[Cid, FakeDownload] = {}
|
|
|
|
async def upload(
|
|
self,
|
|
name: str,
|
|
mime_type: str,
|
|
content: IO,
|
|
timeout: Optional[ClientTimeout] = None,
|
|
) -> Cid:
|
|
data = content.read()
|
|
cid = "Qm" + str(hash(data))
|
|
self.manifests[cid] = Manifest(
|
|
cid=cid,
|
|
treeCid=f"{cid}treeCid",
|
|
datasetSize=len(data),
|
|
mimetype=mime_type,
|
|
blockSize=1,
|
|
filename=name,
|
|
protected=False,
|
|
)
|
|
return cid
|
|
|
|
async def manifest(self, cid: Cid) -> Manifest:
|
|
return self.manifests[cid]
|
|
|
|
async def download(
|
|
self, manifest: Manifest, timeout: Optional[ClientTimeout] = None
|
|
) -> Cid:
|
|
if manifest.treeCid not in self.ongoing_downloads:
|
|
raise ValueError("You need to create a " "download before initiating it")
|
|
return manifest.treeCid
|
|
|
|
def new_download(self, manifest: Manifest) -> FakeDownload:
|
|
"""Creates a download which we can then use to simulate progress."""
|
|
handle = FakeDownload(manifest)
|
|
self.ongoing_downloads[manifest.treeCid] = handle
|
|
return handle
|
|
|
|
async def download_status(self, dataset: Cid) -> DownloadStatus:
|
|
download = self.ongoing_downloads[dataset]
|
|
if download.exception:
|
|
raise download.exception
|
|
return DownloadStatus(
|
|
downloaded=download.downloaded,
|
|
total=download.manifest.block_count,
|
|
)
|