1# Copyright 2016 by Jacek Smietanski.  All rights reserved.
2# This file is part of the Biopython distribution and governed by your
3# choice of the "Biopython License Agreement" or the "BSD 3-Clause License".
4# Please see the LICENSE file that should have been included as part of this
5# package.
6
7"""Testing access to the PDB over the internet."""
8
9import contextlib
10import os
11import shutil
12import tempfile
13import unittest
14
15# We want to test this module:
16from Bio.PDB.PDBList import PDBList
17
18import requires_internet
19
20requires_internet.check()
21
22
23class TestPBDListGetList(unittest.TestCase):
24    """Test methods responsible for getting lists of entries."""
25
26    def test_get_recent_changes(self):
27        """Tests the Bio.PDB.PDBList.get_recent_changes method."""
28        # obsolete_pdb declared to prevent from creating the "obsolete" directory
29        pdblist = PDBList(obsolete_pdb="unimportant")
30        url = pdblist.pdb_server + "/pub/pdb/data/status/latest/added.pdb"
31        entries = pdblist.get_status_list(url)
32        self.assertIsNotNone(entries)
33
34    def test_get_all_entries(self):
35        """Tests the Bio.PDB.PDBList.get_all_entries method."""
36        # obsolete_pdb declared to prevent from creating the "obsolete" directory
37        pdblist = PDBList(obsolete_pdb="unimportant")
38        entries = pdblist.get_all_entries()
39        # As number of entries constantly grow, test checks if a certain number was
40        # exceeded
41        self.assertGreater(len(entries), 100000)
42
43    def test_get_all_obsolete(self):
44        """Tests the Bio.PDB.PDBList.get_all_obsolete method."""
45        # obsolete_pdb declared to prevent from creating the "obsolete" directory
46        pdblist = PDBList(obsolete_pdb="unimportant")
47        entries = pdblist.get_all_obsolete()
48        # As number of obsolete entries constantly grow, test checks if a certain number
49        # was exceeded
50        self.assertGreater(len(entries), 3000)
51
52
53class TestPDBListGetStructure(unittest.TestCase):
54    """Test methods responsible for getting structures."""
55
56    @contextlib.contextmanager
57    def make_temp_directory(self, directory):
58        temp_dir = tempfile.mkdtemp(dir=directory)
59        try:
60            yield temp_dir
61        finally:
62            shutil.rmtree(temp_dir)
63
64    def check(self, structure, filename, file_format, obsolete=False, pdir=None):
65        with self.make_temp_directory(os.getcwd()) as tmp:
66            pdblist = PDBList(pdb=tmp, obsolete_pdb=os.path.join(tmp, "obsolete"))
67            path = os.path.join(tmp, filename)
68            if pdir:
69                pdir = os.path.join(tmp, pdir)
70            pdblist.retrieve_pdb_file(
71                structure, obsolete=obsolete, pdir=pdir, file_format=file_format
72            )
73            self.assertTrue(os.path.isfile(path))
74            os.remove(path)
75
76    def test_retrieve_pdb_file_small_pdb(self):
77        """Tests retrieving the small molecule in pdb format."""
78        structure = "127d"
79        self.check(
80            structure, os.path.join(structure[1:3], "pdb%s.ent" % structure), "pdb"
81        )
82
83    def test_retrieve_pdb_file_large_pdb(self):
84        """Tests retrieving the bundle for large molecule in pdb-like format."""
85        structure = "3k1q"
86        self.check(
87            structure,
88            os.path.join(structure[1:3], "%s-pdb-bundle.tar" % structure),
89            "bundle",
90        )
91
92    def test_retrieve_pdb_file_obsolete_pdb(self):
93        """Tests retrieving the obsolete molecule in pdb format."""
94        structure = "347d"
95        self.check(
96            structure,
97            os.path.join("obsolete", structure[1:3], "pdb%s.ent" % structure),
98            "pdb",
99            obsolete=True,
100        )
101
102    def test_retrieve_pdb_file_obsolete_mmcif(self):
103        """Tests retrieving the obsolete molecule in mmcif format."""
104        structure = "347d"
105        self.check(
106            structure,
107            os.path.join("obsolete", structure[1:3], "%s.cif" % structure),
108            "mmCif",
109            obsolete=True,
110        )
111
112    def test_retrieve_pdb_file_mmcif(self):
113        """Tests retrieving the (non-obsolete) molecule in mmcif format."""
114        structure = "127d"
115        self.check(
116            structure, os.path.join(structure[1:3], "%s.cif" % structure), "mmCif"
117        )
118
119    def test_retrieve_pdb_file_obsolete_xml(self):
120        """Tests retrieving the obsolete molecule in mmcif format."""
121        structure = "347d"
122        self.check(
123            structure,
124            os.path.join("obsolete", structure[1:3], "%s.xml" % structure),
125            "xml",
126            obsolete=True,
127        )
128
129    def test_retrieve_pdb_file_xml(self):
130        """Tests retrieving the (non obsolete) molecule in xml format."""
131        structure = "127d"
132        self.check(structure, os.path.join(structure[1:3], "%s.xml" % structure), "xml")
133
134    def test_retrieve_pdb_file_mmtf(self):
135        """Tests retrieving the molecule in mmtf format."""
136        structure = "127d"
137        self.check(
138            structure, os.path.join(structure[1:3], "%s.mmtf" % structure), "mmtf"
139        )
140
141    def test_double_retrieve(self):
142        """Tests retrieving the same file to different directories."""
143        structure = "127d"
144        self.check(
145            structure, os.path.join("a", "%s.cif" % structure), "mmCif", pdir="a"
146        )
147        self.check(
148            structure, os.path.join("b", "%s.cif" % structure), "mmCif", pdir="b"
149        )
150
151
152if __name__ == "__main__":
153    runner = unittest.TextTestRunner(verbosity=2)
154    unittest.main(testRunner=runner)
155