|
| 1 | +"""Shared stub factories for tests. No unittest.mock anywhere.""" |
| 2 | + |
| 3 | +from datetime import datetime |
| 4 | +from types import SimpleNamespace |
| 5 | + |
| 6 | +from zotero_arxiv_daily.protocol import CorpusPaper, Paper |
| 7 | + |
| 8 | + |
| 9 | +# --------------------------------------------------------------------------- |
| 10 | +# OpenAI client stub |
| 11 | +# --------------------------------------------------------------------------- |
| 12 | + |
| 13 | +_AFFILIATION_MARKER = "You are an assistant who perfectly extracts affiliations" |
| 14 | +_AFFILIATION_RESPONSE = '["TsingHua University","Peking University"]' |
| 15 | +_TLDR_RESPONSE = "Hello! How can I assist you today?" |
| 16 | + |
| 17 | + |
| 18 | +def _make_chat_response(content: str) -> SimpleNamespace: |
| 19 | + return SimpleNamespace( |
| 20 | + choices=[ |
| 21 | + SimpleNamespace( |
| 22 | + message=SimpleNamespace(content=content), |
| 23 | + finish_reason="stop", |
| 24 | + index=0, |
| 25 | + ) |
| 26 | + ], |
| 27 | + id="chatcmpl-stub", |
| 28 | + created=1765197615, |
| 29 | + model="gpt-4o-mini-2024-07-18", |
| 30 | + object="chat.completion", |
| 31 | + ) |
| 32 | + |
| 33 | + |
| 34 | +def _stub_chat_create(**kwargs): |
| 35 | + messages = kwargs.get("messages", []) |
| 36 | + request_str = str(messages) |
| 37 | + if _AFFILIATION_MARKER in request_str: |
| 38 | + return _make_chat_response(_AFFILIATION_RESPONSE) |
| 39 | + return _make_chat_response(_TLDR_RESPONSE) |
| 40 | + |
| 41 | + |
| 42 | +def _stub_embeddings_create(**kwargs): |
| 43 | + inputs = kwargs.get("input", []) |
| 44 | + n = len(inputs) if isinstance(inputs, list) else 1 |
| 45 | + return SimpleNamespace( |
| 46 | + data=[SimpleNamespace(embedding=[0.1, 0.2, 0.3], index=i, object="embedding") for i in range(n)], |
| 47 | + model="text-embedding-3-large", |
| 48 | + object="list", |
| 49 | + ) |
| 50 | + |
| 51 | + |
| 52 | +def make_stub_openai_client(): |
| 53 | + """Return a SimpleNamespace that quacks like openai.OpenAI(). |
| 54 | +
|
| 55 | + chat.completions.create() and embeddings.create() behave identically |
| 56 | + to the Docker mock_openai server that CI previously relied on. |
| 57 | + """ |
| 58 | + return SimpleNamespace( |
| 59 | + chat=SimpleNamespace( |
| 60 | + completions=SimpleNamespace(create=_stub_chat_create), |
| 61 | + ), |
| 62 | + embeddings=SimpleNamespace(create=_stub_embeddings_create), |
| 63 | + ) |
| 64 | + |
| 65 | + |
| 66 | +# --------------------------------------------------------------------------- |
| 67 | +# Zotero client stub |
| 68 | +# --------------------------------------------------------------------------- |
| 69 | + |
| 70 | +_DEFAULT_COLLECTIONS = [ |
| 71 | + { |
| 72 | + "key": "COL1", |
| 73 | + "data": {"name": "survey", "parentCollection": False}, |
| 74 | + }, |
| 75 | + { |
| 76 | + "key": "COL2", |
| 77 | + "data": {"name": "topic-a", "parentCollection": "COL1"}, |
| 78 | + }, |
| 79 | +] |
| 80 | + |
| 81 | +_DEFAULT_ITEMS = [ |
| 82 | + { |
| 83 | + "data": { |
| 84 | + "title": "Stub Paper 1", |
| 85 | + "abstractNote": "Abstract of stub paper 1.", |
| 86 | + "dateAdded": "2026-01-15T10:00:00Z", |
| 87 | + "collections": ["COL2"], |
| 88 | + }, |
| 89 | + }, |
| 90 | + { |
| 91 | + "data": { |
| 92 | + "title": "Stub Paper 2", |
| 93 | + "abstractNote": "Abstract of stub paper 2.", |
| 94 | + "dateAdded": "2026-02-20T12:00:00Z", |
| 95 | + "collections": ["COL1"], |
| 96 | + }, |
| 97 | + }, |
| 98 | +] |
| 99 | + |
| 100 | + |
| 101 | +def make_stub_zotero_client(collections=None, items=None): |
| 102 | + """Return a SimpleNamespace that quacks like pyzotero.zotero.Zotero. |
| 103 | +
|
| 104 | + Supports the call patterns used by Executor.fetch_zotero_corpus(): |
| 105 | + zot.everything(zot.collections()) |
| 106 | + zot.everything(zot.items(itemType=...)) |
| 107 | + """ |
| 108 | + cols = collections if collections is not None else _DEFAULT_COLLECTIONS |
| 109 | + itms = items if items is not None else _DEFAULT_ITEMS |
| 110 | + |
| 111 | + def everything(generator): |
| 112 | + return generator |
| 113 | + |
| 114 | + def collections_fn(): |
| 115 | + return cols |
| 116 | + |
| 117 | + def items_fn(**kwargs): |
| 118 | + return itms |
| 119 | + |
| 120 | + return SimpleNamespace( |
| 121 | + everything=everything, |
| 122 | + collections=collections_fn, |
| 123 | + items=items_fn, |
| 124 | + ) |
| 125 | + |
| 126 | + |
| 127 | +# --------------------------------------------------------------------------- |
| 128 | +# SMTP stub |
| 129 | +# --------------------------------------------------------------------------- |
| 130 | + |
| 131 | + |
| 132 | +def make_stub_smtp(sent_emails: list): |
| 133 | + """Return a class that records calls to sendmail(). |
| 134 | +
|
| 135 | + Usage: |
| 136 | + sent = [] |
| 137 | + monkeypatch.setattr(smtplib, "SMTP", make_stub_smtp(sent)) |
| 138 | + ... |
| 139 | + assert len(sent) == 1 |
| 140 | + sender, recipients, body = sent[0] |
| 141 | + """ |
| 142 | + |
| 143 | + class StubSMTP: |
| 144 | + def __init__(self, *args, **kwargs): |
| 145 | + pass |
| 146 | + |
| 147 | + def starttls(self): |
| 148 | + pass |
| 149 | + |
| 150 | + def login(self, user, password): |
| 151 | + pass |
| 152 | + |
| 153 | + def sendmail(self, sender, recipients, msg): |
| 154 | + sent_emails.append((sender, recipients, msg)) |
| 155 | + |
| 156 | + def quit(self): |
| 157 | + pass |
| 158 | + |
| 159 | + return StubSMTP |
| 160 | + |
| 161 | + |
| 162 | +# --------------------------------------------------------------------------- |
| 163 | +# Paper / CorpusPaper factories |
| 164 | +# --------------------------------------------------------------------------- |
| 165 | + |
| 166 | + |
| 167 | +def make_sample_paper(**overrides) -> Paper: |
| 168 | + defaults = dict( |
| 169 | + source="arxiv", |
| 170 | + title="Sample Paper Title", |
| 171 | + authors=["Author A", "Author B", "Author C"], |
| 172 | + abstract="This paper explores a novel approach to widget engineering.", |
| 173 | + url="https://arxiv.org/abs/2026.00001", |
| 174 | + pdf_url="https://arxiv.org/pdf/2026.00001", |
| 175 | + full_text="\\begin{document} Some text. \\end{document}", |
| 176 | + tldr=None, |
| 177 | + affiliations=None, |
| 178 | + score=None, |
| 179 | + ) |
| 180 | + defaults.update(overrides) |
| 181 | + return Paper(**defaults) |
| 182 | + |
| 183 | + |
| 184 | +def make_sample_corpus(n: int = 3) -> list[CorpusPaper]: |
| 185 | + return [ |
| 186 | + CorpusPaper( |
| 187 | + title=f"Corpus Paper {i}", |
| 188 | + abstract=f"Abstract for corpus paper {i}.", |
| 189 | + added_date=datetime(2026, 1, 1 + i), |
| 190 | + paths=[f"2026/survey/topic-{i}"], |
| 191 | + ) |
| 192 | + for i in range(n) |
| 193 | + ] |
| 194 | + |
| 195 | + |
| 196 | +# --------------------------------------------------------------------------- |
| 197 | +# bioRxiv canned API response |
| 198 | +# --------------------------------------------------------------------------- |
| 199 | + |
| 200 | +SAMPLE_BIORXIV_API_RESPONSE = { |
| 201 | + "messages": [{"status": "ok"}], |
| 202 | + "collection": [ |
| 203 | + { |
| 204 | + "doi": "10.1101/2026.03.01.000001", |
| 205 | + "title": "A biorxiv paper", |
| 206 | + "authors": "Smith, J.; Doe, A.; Lee, K.", |
| 207 | + "abstract": "We present a novel finding.", |
| 208 | + "date": "2026-03-02", |
| 209 | + "category": "bioinformatics", |
| 210 | + "version": "1", |
| 211 | + }, |
| 212 | + { |
| 213 | + "doi": "10.1101/2026.03.01.000002", |
| 214 | + "title": "Another biorxiv paper", |
| 215 | + "authors": "Wang, L.; Chen, M.", |
| 216 | + "abstract": "We replicate a key result.", |
| 217 | + "date": "2026-03-02", |
| 218 | + "category": "genomics", |
| 219 | + "version": "1", |
| 220 | + }, |
| 221 | + { |
| 222 | + "doi": "10.1101/2026.03.01.000003", |
| 223 | + "title": "Old biorxiv paper", |
| 224 | + "authors": "Old, R.", |
| 225 | + "abstract": "Yesterday's paper.", |
| 226 | + "date": "2026-03-01", |
| 227 | + "category": "bioinformatics", |
| 228 | + "version": "1", |
| 229 | + }, |
| 230 | + ], |
| 231 | +} |
0 commit comments