From 6f8a58e7e106c34efc62948bb3bd87e126b21efd Mon Sep 17 00:00:00 2001
From: Logan Ward <ward.logan.t@gmail.com>
Date: Mon, 30 Aug 2021 13:01:38 -0400
Subject: [PATCH] Add support for creating an ADC client

---
 polybot/config.py             | 14 ++++++++++++--
 polybot/sample.py             | 10 +++++++---
 polybot/tests/conftest.py     |  7 -------
 polybot/tests/test_samples.py | 28 +++++++++++++++++++++++++---
 4 files changed, 44 insertions(+), 15 deletions(-)

diff --git a/polybot/config.py b/polybot/config.py
index f7c8f8d..1090fe9 100644
--- a/polybot/config.py
+++ b/polybot/config.py
@@ -3,6 +3,7 @@
 from typing import Optional, Tuple, List
 from urllib.parse import urlparse
 
+from adc.client import ADCClient
 from colmena.redis.queue import ClientQueues, TaskServerQueues
 from pydantic import BaseSettings, Field, HttpUrl, RedisDsn
 
@@ -12,8 +13,9 @@
 class Settings(BaseSettings):
     """Settings for the web service"""
 
-    # Sample handling
-    sample_folder: Path = Field(_run_folder / "samples", description="Path in which to store the samples")
+    # Configuration related to data storage on the Argonne Data Cloud
+    adc_access_token: Optional[str] = Field(None, description='Token for accessing the Argonne Data Cloud')
+    adc_study_id: Optional[str] = Field(None, description='Study ID associated with this experiment')
 
     # Logging
     log_name: Optional[str] = Field(None, description="Name of the log file. If not provided, logs will not be stored")
@@ -66,5 +68,13 @@ def make_server_queue(self) -> TaskServerQueues:
         hostname, port = self.redis_info
         return TaskServerQueues(hostname, port, name='polybot', topics=['robot'] + self.task_queues)
 
+    def generate_adc_client(self) -> ADCClient:
+        """Create an authenticated ADC client
+
+        Returns:
+            A client to the ADC that is ready to make queries
+        """
+        return ADCClient(self.adc_access_token)
+
 
 settings = Settings()
diff --git a/polybot/sample.py b/polybot/sample.py
index 1677738..81b5d41 100644
--- a/polybot/sample.py
+++ b/polybot/sample.py
@@ -9,7 +9,6 @@
 
 from .config import settings
 from .models import Sample
-from adc import sample
 
 
 logger = logging.getLogger(__name__)
@@ -22,8 +21,13 @@ def load_samples() -> Iterator[Sample]:
         Samples in no prescribed order
     """
 
-    sample()
-    for path in settings.sample_folder.glob("*.json"):
+    # Query to get the list of samples in the study
+    adc_client = settings.generate_adc_client()
+    if settings.adc_study_id is None:
+        raise ValueError('The ADC study id is not set. Set your ADC_STUDY_ID environment variable.')
+    study_info = adc_client.get_study(settings.adc_study_id)
+
+    for path in study_info['study']['samples']:
         try:
             yield Sample.parse_file(path)
         except BaseException:
diff --git a/polybot/tests/conftest.py b/polybot/tests/conftest.py
index b629479..1f41caf 100644
--- a/polybot/tests/conftest.py
+++ b/polybot/tests/conftest.py
@@ -13,13 +13,6 @@
 
 @fixture(autouse=True)
 def test_settings():
-    # Redirect the sample folder
-    sample_dir = _test_dir / "test-samples"
-    if sample_dir.is_dir():
-        rmtree(sample_dir)
-    sample_dir.mkdir()
-    settings.sample_folder = sample_dir
-
     # Set up the test Redis service
     settings.redis_url = "rediss://localhost"
 
diff --git a/polybot/tests/test_samples.py b/polybot/tests/test_samples.py
index 3d65837..079e0a8 100644
--- a/polybot/tests/test_samples.py
+++ b/polybot/tests/test_samples.py
@@ -1,10 +1,32 @@
 from polybot.sample import load_samples
-from polybot.config import settings
+
+example_output = {
+    "study": {
+        "id": "U3R1ZHlOb2RlOjU=",
+        "name": "polybot-ai-test",
+        "description": "Test study for developing the Polybot AI planning service.",
+        "keywords": [],
+        "startDate": None,
+        "status": "NEW",
+        "created": "2021-08-30T16:54:41.706854+00:00",
+        "updated": "2021-08-30T16:54:42.127315+00:00",
+        "permissions": [
+            {
+                "user": {
+                    "id": "VXNlck5vZGU6NQ==",
+                    "name": "",
+                    "email": "loganw@uchicago.edu"
+                },
+                "level": "ADMIN"
+            }
+        ],
+        "investigations": [],
+        "samples": []
+    }
+}
 
 
 def test_load(example_sample):
-    with open(settings.sample_folder / "test.json", 'w') as fp:
-        fp.write('Junk')
     samples = list(load_samples())
     assert len(samples) == 1
     assert samples[0].ID == example_sample.ID